1 files changed, 15 insertions, 12 deletions
diff --git a/spec/support/helpers/next_instance_of.rb b/spec/support/helpers/next_instance_of.rb
index 83c788c3d38..a8e9ab2bafe 100644
--- a/spec/support/helpers/next_instance_of.rb
+++ b/spec/support/helpers/next_instance_of.rb
@@ -1,28 +1,31 @@
 # frozen_string_literal: true
 
 module NextInstanceOf
-  def expect_next_instance_of(klass, *new_args)
-    stub_new(expect(klass), *new_args) do |expectation|
-      yield(expectation)
-    end
+  def expect_next_instance_of(klass, *new_args, &blk)
+    stub_new(expect(klass), nil, *new_args, &blk)
   end
 
-  def allow_next_instance_of(klass, *new_args)
-    stub_new(allow(klass), *new_args) do |allowance|
-      yield(allowance)
-    end
+  def expect_next_instances_of(klass, number, *new_args, &blk)
+    stub_new(expect(klass), number, *new_args, &blk)
+  end
+
+  def allow_next_instance_of(klass, *new_args, &blk)
+    stub_new(allow(klass), nil, *new_args, &blk)
+  end
+
+  def allow_next_instances_of(klass, number, *new_args, &blk)
+    stub_new(allow(klass), number, *new_args, &blk)
   end
 
   private
 
-  def stub_new(target, *new_args)
+  def stub_new(target, number, *new_args, &blk)
     receive_new = receive(:new)
+    receive_new.exactly(number).times if number
     receive_new.with(*new_args) if new_args.any?
 
     target.to receive_new.and_wrap_original do |method, *original_args|
-      method.call(*original_args).tap do |instance|
-        yield(instance)
-      end
+      method.call(*original_args).tap(&blk)
     end
   end
 end