Implement ProcStub

src/spectator/dsl/mocks.cr

@@ -308,8 +308,25 @@ module Spectator::DSL
     # allow(dbl).to receive(:foo).and_return(42)
     # expect(dbl.foo).to eq(42)
     # ```
-    macro receive(method, *, _file = __FILE__, _line = __LINE__)
-      ::Spectator::NullStub.new({{method.id.symbolize}}, location: ::Spectator::Location.new({{_file}}, {{_line}}))
+    #
+    # A block can be provided to be run every time the stub is invoked.
+    # The value returned by the block is returned by the stubbed method.
+    #
+    # ```
+    # dbl = dbl(:foobar)
+    # allow(dbl).to receive(:foo) { 42 }
+    # expect(dbl.foo).to eq(42)
+    # ```
+    macro receive(method, *, _file = __FILE__, _line = __LINE__, &block)
+      {% if block %}
+        %proc = ->(%args : ::Spectator::AbstractArguments) {
+          {% if !block.args.empty? %}{{*block.args}} = %args {% end %}
+          {{block.body}}
+        }
+        ::Spectator::ProcStub.new({{method.id.symbolize}}, %proc, location: ::Spectator::Location.new({{_file}}, {{_line}}))
+      {% else %}
+        ::Spectator::NullStub.new({{method.id.symbolize}}, location: ::Spectator::Location.new({{_file}}, {{_line}}))
+      {% end %}
     end
   end
 end

src/spectator/mocks/proc_stub.cr

@@ -0,0 +1,23 @@
+require "../location"
+require "./arguments"
+require "./typed_stub"
+
+module Spectator
+  # Stub that responds with a value returned by calling a proc.
+  class ProcStub(T) < TypedStub(T)
+    # Invokes the stubbed implementation.
+    def call(call : MethodCall) : T
+      @proc.call(call.arguments)
+    end
+
+    # Creates the stub.
+    def initialize(method : Symbol, @proc : Proc(AbstractArguments, T), constraint : AbstractArguments? = nil, location : Location? = nil)
+      super(method, constraint, location)
+    end
+
+    # Creates the stub.
+    def initialize(method : Symbol, constraint : AbstractArguments? = nil, location : Location? = nil, &block : Proc(AbstractArguments, T))
+      initialize(method, block, constraint, location)
+    end
+  end
+end