From 27acf7bced1c3f859cd4859fc6a3793a75677dd9 Mon Sep 17 00:00:00 2001
From: Michael Miller <icy.arctic.fox@gmail.com>
Date: Tue, 19 Nov 2019 17:28:00 -0700
Subject: [PATCH] Allow receive block syntax

---
 src/spectator/dsl/mocks.cr              | 8 ++++++--
 src/spectator/mocks/proc_method_stub.cr | 4 ++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/spectator/dsl/mocks.cr b/src/spectator/dsl/mocks.cr
index 7470f16..ef24d6c 100644
--- a/src/spectator/dsl/mocks.cr
+++ b/src/spectator/dsl/mocks.cr
@@ -126,9 +126,13 @@ module Spectator::DSL
     ::Spectator::Mocks::ExpectAnyInstance({{type}}).new(%source)
   end
 
-  macro receive(method_name, _source_file = __FILE__, _source_line = __LINE__)
+  macro receive(method_name, _source_file = __FILE__, _source_line = __LINE__, &block)
     %source = ::Spectator::Source.new({{_source_file}}, {{_source_line}})
-    ::Spectator::Mocks::NilMethodStub.new({{method_name.id.symbolize}}, %source)
+    {% if block.is_a?(Nop) %}
+      ::Spectator::Mocks::NilMethodStub.new({{method_name.id.symbolize}}, %source)
+    {% else %}
+      ::Spectator::Mocks::ProcMethodStub.create({{method_name.id.symbolize}}, %source) { {{block.body}} }
+    {% end %}
   end
 
   macro receive_messages(_source_file = __FILE__, _source_line = __LINE__, **stubs)
diff --git a/src/spectator/mocks/proc_method_stub.cr b/src/spectator/mocks/proc_method_stub.cr
index b0d921f..eeb0569 100644
--- a/src/spectator/mocks/proc_method_stub.cr
+++ b/src/spectator/mocks/proc_method_stub.cr
@@ -7,6 +7,10 @@ module Spectator::Mocks
       super(name, source, args)
     end
 
+    def self.create(name, source, args = nil, &block : -> T) forall T
+      ProcMethodStub.new(name, source, block, args)
+    end
+
     def call(_args : GenericArguments(T2, NT2), rt : RT.class) forall T2, NT2, RT
       @proc.call
     end