From 5c910e5a85fa3342f7587de4d0a4e49e10c567f6 Mon Sep 17 00:00:00 2001
From: Michael Miller <icy.arctic.fox@gmail.com>
Date: Sun, 9 Oct 2022 13:57:28 -0600
Subject: [PATCH] Clear stubs defined with `expect().to receive()` syntax after
 test finishes

---
 CHANGELOG.md                 |  3 +++
 spec/docs/mocks_spec.cr      |  4 ++--
 src/spectator/expectation.cr | 24 ++++++++++++++++++++----
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 133aa2d..94f89bf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Fixed
+- Clear stubs defined with `expect().to receive()` syntax after test finishes to prevent leakage between tests.
+
 ### Removed
 - Removed support for stubbing undefined (untyped) methods in lazy doubles. Avoids possible segfault.
 
diff --git a/spec/docs/mocks_spec.cr b/spec/docs/mocks_spec.cr
index 6e164ae..5908787 100644
--- a/spec/docs/mocks_spec.cr
+++ b/spec/docs/mocks_spec.cr
@@ -146,9 +146,9 @@ Spectator.describe "Mocks Docs" do
       inst.something
     end
 
-    it "leaks stubs to other examples" do
+    it "reverts to default stub for other examples" do
       inst = mock(MyStruct)
-      expect(inst.something).to eq(7) # Previous stub was leaked.
+      expect(inst.something).to eq(5) # Default stub used instead of original behavior.
     end
   end
 end
diff --git a/src/spectator/expectation.cr b/src/spectator/expectation.cr
index bfc0248..8028429 100644
--- a/src/spectator/expectation.cr
+++ b/src/spectator/expectation.cr
@@ -160,9 +160,17 @@ module Spectator
           stubbable._spectator_define_stub(unconstrained_stub)
         end
 
+        # Apply the stub that is expected to be called.
         stubbable._spectator_define_stub(stub)
-        matcher = Matchers::ReceiveMatcher.new(stub)
-        to_eventually(matcher, message)
+
+        # Check if the stub was invoked after the test completes.
+        Harness.current.defer do
+          matcher = Matchers::ReceiveMatcher.new(stub)
+          to(matcher, message)
+        ensure
+          # Prevent leaking stubs between tests.
+          stubbable._spectator_remove_stub(stub)
+        end
       end
 
       # Asserts that some criteria defined by the matcher is eventually satisfied.
@@ -190,9 +198,17 @@ module Spectator
           stubbable._spectator_define_stub(unconstrained_stub)
         end
 
+        # Apply the stub that could be called in case it is.
         stubbable._spectator_define_stub(stub)
-        matcher = Matchers::ReceiveMatcher.new(stub)
-        to_never(matcher, message)
+
+        # Check if the stub was invoked after the test completes.
+        Harness.current.defer do
+          matcher = Matchers::ReceiveMatcher.new(stub)
+          to_not(matcher, message)
+        ensure
+          # Prevent leaking stubs between tests.
+          stubbable._spectator_remove_stub(stub)
+        end
       end
 
       # :ditto: