From 2ddd80cf1dc17a6e1e23a5884d54390e736f55e4 Mon Sep 17 00:00:00 2001
From: Yuri Pereira Constante <ypconstante@gmail.com>
Date: Sat, 22 Jun 2024 10:26:24 -0300
Subject: [PATCH] Remove redundant ownership check

---
 lib/mox.ex | 80 +++++++++++++++++++++++++++---------------------------
 1 file changed, 40 insertions(+), 40 deletions(-)

diff --git a/lib/mox.ex b/lib/mox.ex
index a211814..427e9dd 100644
--- a/lib/mox.ex
+++ b/lib/mox.ex
@@ -685,8 +685,11 @@ defmodule Mox do
       :ok ->
         :ok
 
+      {:error, error} when is_binary(error) ->
+        raise ArgumentError, error
+
       {:error, error} ->
-        raise ArgumentError, expectation_error_to_message(error, mock)
+        raise error
     end
   end
 
@@ -697,31 +700,14 @@ defmodule Mox do
       :ok ->
         :ok
 
+      {:error, error} when is_binary(error) ->
+        raise ArgumentError, error
+
       {:error, error} ->
-        raise ArgumentError, expectation_error_to_message(error, mock)
+        raise error
     end
   end
 
-  defp expectation_error_to_message({:currently_allowed, owner_pid}, mock) do
-    inspected = inspect(self())
-
-    """
-    cannot add expectations/stubs to #{inspect(mock)} in the current process (#{inspected}) \
-    because the process has been allowed by #{inspect(owner_pid)}. \
-    You cannot define expectations/stubs in a process that has been allowed
-    """
-  end
-
-  defp expectation_error_to_message({:not_shared_owner, global_pid}, mock) do
-    inspected = inspect(self())
-
-    """
-    cannot add expectations/stubs to #{inspect(mock)} in the current process (#{inspected}) \
-    because Mox is in global mode and the global process is #{inspect(global_pid)}. \
-    Only the process that set Mox to global can set expectations/stubs in global mode
-    """
-  end
-
   @doc """
   Allows other processes to share expectations and stubs
   defined by owner process.
@@ -948,13 +934,34 @@ defmodule Mox do
   end
 
   defp add_expectations(owner_pid, mock, key_expectation_list) do
-    case ensure_pid_can_add_expectation(owner_pid, mock) do
-      :ok ->
-        update_fun = &{:ok, init_or_merge_expectations(&1, key_expectation_list)}
-        :ok = get_and_update!(owner_pid, mock, update_fun)
+    update_fun = &{:ok, init_or_merge_expectations(&1, key_expectation_list)}
+
+    case get_and_update(owner_pid, mock, update_fun) do
+      {:ok, value} ->
+        value
 
-      {:error, reason} ->
-        {:error, reason}
+      {:error, %NimbleOwnership.Error{reason: {:already_allowed, _}}} ->
+        inspected = inspect(self())
+
+        {:error,
+         """
+         cannot add expectations/stubs to #{inspect(mock)} in the current process (#{inspected}) \
+         because the process has been allowed by #{inspect(owner_pid)}. \
+         You cannot define expectations/stubs in a process that has been allowed
+         """}
+
+      {:error, %NimbleOwnership.Error{reason: {:not_shared_owner, global_pid}}} ->
+        inspected = inspect(self())
+
+        {:error,
+         """
+         cannot add expectations/stubs to #{inspect(mock)} in the current process (#{inspected}) \
+         because Mox is in global mode and the global process is #{inspect(global_pid)}. \
+         Only the process that set Mox to global can set expectations/stubs in global mode
+         """}
+
+      {:error, error} ->
+        {:error, error}
     end
   end
 
@@ -981,17 +988,6 @@ defmodule Mox do
     end
   end
 
-  # Make sure that the owner_pid is either the owner or that the mock
-  # isn't owned yet.
-  defp ensure_pid_can_add_expectation(owner_pid, mock) do
-    case NimbleOwnership.fetch_owner(@this, [owner_pid], mock, @timeout) do
-      :error -> :ok
-      {tag, ^owner_pid} when tag in [:ok, :shared_owner] -> :ok
-      {:shared_owner, other_owner} -> {:error, {:not_shared_owner, other_owner}}
-      {:ok, other_owner} -> {:error, {:currently_allowed, other_owner}}
-    end
-  end
-
   defp fetch_owner_from_callers(caller_pids, mock) do
     # If the mock doesn't have an owner, it can't have expectations so we return :no_expectation.
     case NimbleOwnership.fetch_owner(@this, caller_pids, mock, @timeout) do
@@ -1000,8 +996,12 @@ defmodule Mox do
     end
   end
 
+  defp get_and_update(owner_pid, mock, update_fun) do
+    NimbleOwnership.get_and_update(@this, owner_pid, mock, update_fun, @timeout)
+  end
+
   defp get_and_update!(owner_pid, mock, update_fun) do
-    case NimbleOwnership.get_and_update(@this, owner_pid, mock, update_fun, @timeout) do
+    case get_and_update(owner_pid, mock, update_fun) do
       {:ok, return} -> return
       {:error, %NimbleOwnership.Error{} = error} -> raise error
     end