From 2667b2d0a4e5da24529752a91b3d75c7f77c4cc1 Mon Sep 17 00:00:00 2001
From: Tony van Riet <306897+tonyvanriet@users.noreply.github.com>
Date: Sun, 18 Aug 2019 02:08:34 -0500
Subject: [PATCH] Always overwrite the stub when merging new expects or stubs
 (#70)

---
 lib/mox.ex        |  5 +++++
 lib/mox/server.ex |  4 ++--
 test/mox_test.exs | 15 +++++++++++++--
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/lib/mox.ex b/lib/mox.ex
index 4c8b875..264d50f 100644
--- a/lib/mox.ex
+++ b/lib/mox.ex
@@ -359,6 +359,11 @@ defmodule Mox do
   `expect/4` can also be invoked multiple times for the same
   name/arity, allowing you to give different behaviours on each
   invocation.
+
+  When `expect/4` is invoked, any previously declared `stub` for the same `name` and arity will
+  be removed. This ensures that `expect` will fail if the function is called more than `n` times.
+  If a `stub/3` is invoked **after** `expect/4` for the same `name` and arity, the stub will be
+  used after all expectations are fulfilled.
   """
   def expect(mock, name, n \\ 1, code)
       when is_atom(mock) and is_atom(name) and is_integer(n) and n >= 0 and is_function(code) do
diff --git a/lib/mox/server.ex b/lib/mox/server.ex
index 4841446..ae17010 100644
--- a/lib/mox/server.ex
+++ b/lib/mox/server.ex
@@ -240,7 +240,7 @@ defmodule Mox.Server do
     end
   end
 
-  defp merge_expectation({current_n, current_calls, current_stub}, {n, calls, stub}) do
-    {current_n + n, current_calls ++ calls, stub || current_stub}
+  defp merge_expectation({current_n, current_calls, _current_stub}, {n, calls, stub}) do
+    {current_n + n, current_calls ++ calls, stub}
   end
 end
diff --git a/test/mox_test.exs b/test/mox_test.exs
index 0fb51cc..f592edd 100644
--- a/test/mox_test.exs
+++ b/test/mox_test.exs
@@ -292,6 +292,17 @@ defmodule MoxTest do
       end
     end
 
+    test "raises if all expectations are consumed, even when a stub is defined" do
+      stub(CalcMock, :add, fn _, _ -> :stub end)
+
+      expect(CalcMock, :add, 1, fn _, _ -> :expected end)
+      assert CalcMock.add(2, 3) == :expected
+
+      assert_raise Mox.UnexpectedCallError, fn ->
+        CalcMock.add(2, 3)
+      end
+    end
+
     test "raises if you try to add expectations from non global process" do
       set_mox_global()
 
@@ -485,11 +496,11 @@ defmodule MoxTest do
       end)
     end
 
-    test "invokes stub after expectations are fulfilled" do
+    test "a stub declared after an expect is invoked after all expectations are fulfilled" do
       in_all_modes(fn ->
         CalcMock
-        |> stub(:add, fn _x, _y -> :stub end)
         |> expect(:add, 2, fn _, _ -> :expected end)
+        |> stub(:add, fn _x, _y -> :stub end)
 
         assert CalcMock.add(1, 1) == :expected
         assert CalcMock.add(1, 1) == :expected