From 86b49dc20e1fdf9869b956a737b417444d3803c5 Mon Sep 17 00:00:00 2001
From: Michael Miller <icy.arctic.fox@gmail.com>
Date: Sun, 15 May 2022 00:36:29 -0600
Subject: [PATCH] Add mock registry fetch method

Solves the issue of pre-populating a mock with stubs.
---
 .../mocks/reference_mock_registry_spec.cr     | 24 ++++++++++++++++---
 src/spectator/mocks/mock.cr                   | 10 +++++++-
 .../mocks/reference_mock_registry.cr          | 11 +++++++++
 3 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/spec/spectator/mocks/reference_mock_registry_spec.cr b/spec/spectator/mocks/reference_mock_registry_spec.cr
index a607d6f..1d34f87 100644
--- a/spec/spectator/mocks/reference_mock_registry_spec.cr
+++ b/spec/spectator/mocks/reference_mock_registry_spec.cr
@@ -2,16 +2,16 @@ require "../../spec_helper"
 
 Spectator.describe Spectator::ReferenceMockRegistry do
   subject(registry) { described_class.new }
+  let(obj) { "foobar" }
   let(stub) { Spectator::ValueStub.new(:test, 42) }
+  let(no_stubs) { [] of Spectator::Stub }
 
   it "initially has no stubs" do
-    obj = "foobar"
     expect(registry[obj]).to be_empty
   end
 
   it "stores stubs for an object" do
-    obj = "foobar"
-    expect { registry[obj] << stub }.to change { registry[obj] }.from([] of Spectator::Stub).to([stub])
+    expect { registry[obj] << stub }.to change { registry[obj] }.from(no_stubs).to([stub])
   end
 
   it "isolates stubs between different objects" do
@@ -20,4 +20,22 @@ Spectator.describe Spectator::ReferenceMockRegistry do
     registry[obj2] << Spectator::ValueStub.new(:obj2, 42)
     expect { registry[obj1] << stub }.to_not change { registry[obj2] }
   end
+
+  describe "#fetch" do
+    it "retrieves existing stubs" do
+      registry[obj] << stub
+      expect(registry.fetch(obj) { no_stubs }).to eq([stub])
+    end
+
+    it "stores stubs on the first retrieval" do
+      expect(registry.fetch(obj) { [stub] of Spectator::Stub }).to eq([stub])
+    end
+
+    it "isolates stubs between different objects" do
+      obj1 = "foo"
+      obj2 = "bar"
+      registry[obj2] << Spectator::ValueStub.new(:obj2, 42)
+      expect { registry.fetch(obj2) { no_stubs } }.to_not change { registry[obj2] }
+    end
+  end
 end
diff --git a/src/spectator/mocks/mock.cr b/src/spectator/mocks/mock.cr
index 73694e6..b7613a0 100644
--- a/src/spectator/mocks/mock.cr
+++ b/src/spectator/mocks/mock.cr
@@ -59,7 +59,15 @@ module Spectator
           {% end %}
 
           private def _spectator_stubs
-            @@_spectator_mock_registry[self]
+            @@_spectator_mock_registry.fetch(self) do
+              {% begin %}
+                [
+                  {% for key, value in value_methods %}
+                    ::Spectator::ValueStub.new({{key.id.symbolize}}, {{value}}),
+                  {% end %}
+                ] of ::Spectator::Stub
+              {% end %}
+            end
           end
 
           # Returns the mock's name formatted for user output.
diff --git a/src/spectator/mocks/reference_mock_registry.cr b/src/spectator/mocks/reference_mock_registry.cr
index e6272db..6116e73 100644
--- a/src/spectator/mocks/reference_mock_registry.cr
+++ b/src/spectator/mocks/reference_mock_registry.cr
@@ -23,5 +23,16 @@ module Spectator
       key = Box.box(object)
       @object_stubs[key]
     end
+
+    # Retrieves all stubs defined for a mocked object.
+    #
+    # Yields to the block on the first retrieval.
+    # This allows a mock to populate the registry with initial stubs.
+    def fetch(object : Reference, & : -> Array(Stub))
+      key = Box.box(object)
+      @object_stubs.fetch(key) do
+        @object_stubs[key] = yield
+      end
+    end
   end
 end