Add latest changes from gitlab-org/gitlab@13-4-stable-ee

1 files changed, 33 insertions, 0 deletions

diff --git a/spec/support/helpers/next_found_instance_of.rb b/spec/support/helpers/next_found_instance_of.rb
new file mode 100644
index 00000000000..ff34fcdd1d3
--- /dev/null
+++ b/spec/support/helpers/next_found_instance_of.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module NextFoundInstanceOf
+  ERROR_MESSAGE = 'NextFoundInstanceOf mock helpers can only be used with ActiveRecord targets'
+
+  def expect_next_found_instance_of(klass)
+    check_if_active_record!(klass)
+
+    stub_allocate(expect(klass)) do |expectation|
+      yield(expectation)
+    end
+  end
+
+  def allow_next_found_instance_of(klass)
+    check_if_active_record!(klass)
+
+    stub_allocate(allow(klass)) do |allowance|
+      yield(allowance)
+    end
+  end
+
+  private
+
+  def check_if_active_record!(klass)
+    raise ArgumentError.new(ERROR_MESSAGE) unless klass < ActiveRecord::Base
+  end
+
+  def stub_allocate(target)
+    target.to receive(:allocate).and_wrap_original do |method|
+      method.call.tap { |allocation| yield(allocation) }
+    end
+  end
+end