From 9f54a9e542b46e75b74ac3989ae1b0f8b09c8043 Mon Sep 17 00:00:00 2001
From: Michael Miller <icy.arctic.fox@gmail.com>
Date: Sat, 17 Dec 2022 19:16:38 -0700
Subject: [PATCH] Additional handling for passing blocks

---
 CHANGELOG.md                     |  1 +
 src/spectator/mocks/stubbable.cr | 62 ++++++++++++++++++++++----------
 2 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e700a80..645a115 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 ### Fixed
 - Fix macro logic to support free variables, 'self', and variants on stubbed methods. [#48](https://github.com/icy-arctic-fox/spectator/issues/48)
+- Fix method stubs used on methods that capture blocks.
 
 ### Changed
 - Simplify string representation of mock-related types.
diff --git a/src/spectator/mocks/stubbable.cr b/src/spectator/mocks/stubbable.cr
index 6ded8ee..9c32c47 100644
--- a/src/spectator/mocks/stubbable.cr
+++ b/src/spectator/mocks/stubbable.cr
@@ -130,14 +130,25 @@ module Spectator
          # Workaround for Crystal not propagating block with previous_def/super.
          if method.accepts_block?
            original += "("
-           method.args.each_with_index do |arg, i|
-             original += '*' if method.splat_index == i
-             original += arg.name.stringify
-             original += ", "
-           end
-           if method.double_splat
-             original += method.double_splat.stringify
-             original += ", "
+           if method.splat_index
+             method.args.each_with_index do |arg, i|
+               if i == method.splat_index
+                 original += '*'
+                 if arg.internal_name && arg.internal_name.size > 0
+                   original += "#{arg.internal_name}, "
+                 end
+                 original += "**#{method.double_splat}, " if method.double_splat
+               elsif i > method.splat_index
+                 original += "#{arg.name}: #{arg.internal_name}"
+               else
+                 original += "#{arg.internal_name}, "
+               end
+             end
+           else
+             method.args.each do |arg|
+               original += "#{arg.internal_name}, "
+             end
+             original += "**#{method.double_splat}, " if method.double_splat
            end
            # If the block is captured (i.e. `&block` syntax), it must be passed along as an argument.
            # Otherwise, use `yield` to forward the block.
@@ -269,14 +280,25 @@ module Spectator
            # Workaround for Crystal not propagating block with previous_def/super.
            if method.accepts_block?
              original += "("
-             method.args.each_with_index do |arg, i|
-               original += '*' if method.splat_index == i
-               original += arg.name.stringify
-               original += ", "
-             end
-             if method.double_splat
-               original += method.double_splat.stringify
-               original += ", "
+             if method.splat_index
+               method.args.each_with_index do |arg, i|
+                 if i == method.splat_index
+                   original += '*'
+                   if arg.internal_name && arg.internal_name.size > 0
+                     original += "#{arg.internal_name}, "
+                   end
+                   original += "**#{method.double_splat}, " if method.double_splat
+                 elsif i > method.splat_index
+                   original += "#{arg.name}: #{arg.internal_name}"
+                 else
+                   original += "#{arg.internal_name}, "
+                 end
+               end
+             else
+               method.args.each do |arg|
+                 original += "#{arg.internal_name}, "
+               end
+               original += "**#{method.double_splat}, " if method.double_splat
              end
              # If the block is captured (i.e. `&block` syntax), it must be passed along as an argument.
              # Otherwise, use `yield` to forward the block.
@@ -467,9 +489,11 @@ module Spectator
           {% if method.block_arg %}&{{method.block_arg}}{% elsif method.accepts_block? %}&{% end %}
         ){% if method.return_type %} : {{method.return_type}}{% end %}{% if !method.free_vars.empty? %} forall {{method.free_vars.splat}}{% end %}
         {% unless method.abstract? %}
-            {{scope}}{% if method.accepts_block? %}(
-              {% for arg, i in method.args %}{% if i == method.splat_index %}*{% end %}{{arg.name}}, {% end %}
-              {% if method.double_splat %}**{{method.double_splat}}, {% end %}
+            {{scope}}{% if method.accepts_block? %}({% for arg, i in method.args %}
+              {% if i == method.splat_index && arg.internal_name && arg.internal_name.size > 0 %}*{{arg.internal_name}}, {% if method.double_splat %}**{{method.double_splat}}, {% end %}{% end %}
+              {% if method.splat_index && i > method.splat_index %}{{arg.name}}: {{arg.internal_name}}, {% end %}
+              {% if !method.splat_index || i < method.splat_index %}{{arg.internal_name}}, {% end %}{% end %}
+              {% if !method.splat_index && method.double_splat %}**{{method.double_splat}}, {% end %}
               {% captured_block = if method.block_arg && method.block_arg.name && method.block_arg.name.size > 0
                                     method.block_arg.name
                                   else