From 0e556c3d55aa8e74f7d77304c8c7eb2d77a9e41a Mon Sep 17 00:00:00 2001
From: Michael Miller <icy.arctic.fox@gmail.com>
Date: Tue, 12 Jul 2022 20:40:27 -0600
Subject: [PATCH] Clear stubs and recorded calls after example completes

---
 src/spectator/dsl/mocks.cr          | 38 ++++++++++++++++++++---------
 src/spectator/harness.cr            | 20 +++++++++++++++
 src/spectator/mocks/stubbable.cr    |  6 +++++
 src/spectator/mocks/stubbed_type.cr |  6 +++++
 4 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/src/spectator/dsl/mocks.cr b/src/spectator/dsl/mocks.cr
index 1796be6..df016ae 100644
--- a/src/spectator/dsl/mocks.cr
+++ b/src/spectator/dsl/mocks.cr
@@ -94,11 +94,15 @@ module Spectator::DSL
  end
  found_tuple = found_tuples.last %}
 
-      {% if found_tuple %}
-        {{found_tuple[2].id}}.new({{**value_methods}})
-      {% else %}
-        ::Spectator::LazyDouble.new({{name}}, {{**value_methods}})
-      {% end %}
+      begin
+        %double = {% if found_tuple %}
+                    {{found_tuple[2].id}}.new({{**value_methods}})
+                  {% else %}
+                    ::Spectator::LazyDouble.new({{name}}, {{**value_methods}})
+                  {% end %}
+        ::Spectator::Harness.current?.try(&.cleanup { %double._spectator_reset })
+        %double
+      end
     end
 
     # Instantiates a class double.
@@ -148,11 +152,15 @@ module Spectator::DSL
  end
  found_tuple = found_tuples.last %}
 
-      {% if found_tuple %}
-        {{found_tuple[2].id}}
-      {% else %}
-        ::Spectator::LazyDouble
-      {% end %}
+      begin
+        %double = {% if found_tuple %}
+                    {{found_tuple[2].id}}
+                  {% else %}
+                    ::Spectator::LazyDouble
+                  {% end %}
+        ::Spectator::Harness.current?.try(&.cleanup { %double._spectator_reset })
+        %double
+      end
     end
 
     # Defines or instantiates a double.
@@ -286,11 +294,14 @@ module Spectator::DSL
  found_tuple = found_tuples.last %}
 
       {% if found_tuple %}
-        {{found_tuple[2].id}}.new.tap do |%mock|
+        begin
+          %mock = {{found_tuple[2].id}}.new
           {% for key, value in value_methods %}
             %stub{key} = ::Spectator::ValueStub.new({{key.id.symbolize}}, {{value}})
             %mock._spectator_define_stub(%stub{key})
           {% end %}
+          ::Spectator::Harness.current?.try(&.cleanup { %mock._spectator_reset })
+          %mock
         end
       {% else %}
         {% raise "Type `#{type.id}` must be previously mocked before attempting to instantiate." %}
@@ -362,11 +373,14 @@ module Spectator::DSL
  found_tuple = found_tuples.last %}
 
       {% if found_tuple %}
-        {{found_tuple[2].id}}.tap do |%mock|
+        begin
+          %mock = {{found_tuple[2].id}}
           {% for key, value in value_methods %}
             %stub{key} = ::Spectator::ValueStub.new({{key.id.symbolize}}, {{value}})
             %mock._spectator_define_stub(%stub{key})
           {% end %}
+          ::Spectator::Harness.current?.try(&.cleanup { %mock._spectator_reset })
+          %mock
         end
       {% else %}
         {% raise "Type `#{type.id}` must be previously mocked before attempting to instantiate." %}
diff --git a/src/spectator/harness.cr b/src/spectator/harness.cr
index ec21d51..6be48b9 100644
--- a/src/spectator/harness.cr
+++ b/src/spectator/harness.cr
@@ -64,6 +64,7 @@ module Spectator
     end
 
     @deferred = Deque(->).new
+    @cleanup = Deque(->).new
     @expectations = [] of Expectation
     @aggregate : Array(Expectation)? = nil
 
@@ -72,6 +73,7 @@ module Spectator
     def run : Result
       elapsed, error = capture { yield }
       elapsed2, error2 = capture { run_deferred }
+      run_cleanup
       translate(elapsed + elapsed2, error || error2)
     end
 
@@ -97,6 +99,13 @@ module Spectator
       @deferred << block
     end
 
+    # Stores a block of code to be executed at cleanup.
+    # Cleanup is run after everything else, even deferred blocks.
+    # Each cleanup step is wrapped in error handling so that one failure doesn't block the next ones.
+    def cleanup(&block) : Nil
+      @cleanup << block
+    end
+
     def aggregate_failures(label = nil)
       previous = @aggregate
       @aggregate = aggregate = [] of Expectation
@@ -168,5 +177,16 @@ module Spectator
       Log.debug { "Running deferred operations" }
       @deferred.each(&.call)
     end
+
+    # Invokes all cleanup callbacks.
+    # Each callback is wrapped with error handling.
+    private def run_cleanup
+      Log.debug { "Running cleanup" }
+      @cleanup.each do |callback|
+        callback.call
+      rescue e
+        Log.error(exception: e) { "Encountered error during cleanup" }
+      end
+    end
   end
 end
diff --git a/src/spectator/mocks/stubbable.cr b/src/spectator/mocks/stubbable.cr
index 75575cd..d998246 100644
--- a/src/spectator/mocks/stubbable.cr
+++ b/src/spectator/mocks/stubbable.cr
@@ -79,6 +79,12 @@ module Spectator
     # Utility method returning the stubbed type's name formatted for user output.
     abstract def _spectator_stubbed_name : String
 
+    # Clears all previously defined calls and stubs.
+    def _spectator_reset : Nil
+      _spectator_clear_calls
+      _spectator_clear_stubs
+    end
+
     # Redefines a method to accept stubs and provides a default response.
     #
     # The *method* must be a `Def`.
diff --git a/src/spectator/mocks/stubbed_type.cr b/src/spectator/mocks/stubbed_type.cr
index c08bbfd..a5588ee 100644
--- a/src/spectator/mocks/stubbed_type.cr
+++ b/src/spectator/mocks/stubbed_type.cr
@@ -32,6 +32,12 @@ module Spectator
       _spectator_calls.clear
     end
 
+    # Clears all previously defined calls and stubs.
+    def _spectator_reset : Nil
+      _spectator_clear_calls
+      _spectator_clear_stubs
+    end
+
     def _spectator_stub_fallback(call : MethodCall, &)
       Log.trace { "Fallback for #{call} - call original" }
       yield