Implement `expect_next_instance_of` and use it

We need this because `expect_any_instance_of` doesn't work on
prepended models. Now we could use the same code between CE/EE

7 files changed, 30 insertions, 11 deletions

diff --git a/spec/lib/gitlab/bitbucket_import/project_creator_spec.rb b/spec/lib/gitlab/bitbucket_import/project_creator_spec.rb
index c3f528dd6fc..ed6fa3d229f 100644
--- a/spec/lib/gitlab/bitbucket_import/project_creator_spec.rb
+++ b/spec/lib/gitlab/bitbucket_import/project_creator_spec.rb
@@ -25,7 +25,9 @@ describe Gitlab::BitbucketImport::ProjectCreator do
   end
 
   it 'creates project' do
-    allow_any_instance_of(Project).to receive(:add_import_job)
+    expect_next_instance_of(Project) do |project|
+      expect(project).to receive(:add_import_job)
+    end
 
     project_creator = described_class.new(repo, 'vim', namespace, user, access_params)
     project = project_creator.execute
diff --git a/spec/lib/gitlab/gitlab_import/project_creator_spec.rb b/spec/lib/gitlab/gitlab_import/project_creator_spec.rb
index 5ea086e4abd..b814f5fc76c 100644
--- a/spec/lib/gitlab/gitlab_import/project_creator_spec.rb
+++ b/spec/lib/gitlab/gitlab_import/project_creator_spec.rb
@@ -21,7 +21,9 @@ describe Gitlab::GitlabImport::ProjectCreator do
   end
 
   it 'creates project' do
-    allow_any_instance_of(Project).to receive(:add_import_job)
+    expect_next_instance_of(Project) do |project|
+      expect(project).to receive(:add_import_job)
+    end
 
     project_creator = described_class.new(repo, namespace, user, access_params)
     project = project_creator.execute
diff --git a/spec/lib/gitlab/google_code_import/project_creator_spec.rb b/spec/lib/gitlab/google_code_import/project_creator_spec.rb
index 24cd518c77b..b959e006292 100644
--- a/spec/lib/gitlab/google_code_import/project_creator_spec.rb
+++ b/spec/lib/gitlab/google_code_import/project_creator_spec.rb
@@ -16,7 +16,9 @@ describe Gitlab::GoogleCodeImport::ProjectCreator do
   end
 
   it 'creates project' do
-    allow_any_instance_of(Project).to receive(:add_import_job)
+    expect_next_instance_of(Project) do |project|
+      expect(project).to receive(:add_import_job)
+    end
 
     project_creator = described_class.new(repo, namespace, user)
     project = project_creator.execute
diff --git a/spec/lib/gitlab/legacy_github_import/project_creator_spec.rb b/spec/lib/gitlab/legacy_github_import/project_creator_spec.rb
index 972b17d5b12..3d4240fa4ba 100644
--- a/spec/lib/gitlab/legacy_github_import/project_creator_spec.rb
+++ b/spec/lib/gitlab/legacy_github_import/project_creator_spec.rb
@@ -17,7 +17,10 @@ describe Gitlab::LegacyGithubImport::ProjectCreator do
 
   before do
     namespace.add_owner(user)
-    allow_any_instance_of(Project).to receive(:add_import_job)
+
+    expect_next_instance_of(Project) do |project|
+      expect(project).to receive(:add_import_job)
+    end
   end
 
   describe '#execute' do
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index dac609e2545..fdce8e84620 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -69,6 +69,7 @@ RSpec.configure do |config|
   config.include StubFeatureFlags
   config.include StubGitlabCalls
   config.include StubGitlabData
+  config.include ExpectNextInstanceOf
   config.include TestEnv
   config.include Devise::Test::ControllerHelpers, type: :controller
   config.include Devise::Test::IntegrationHelpers, type: :feature
diff --git a/spec/support/helpers/expect_next_instance_of.rb b/spec/support/helpers/expect_next_instance_of.rb
new file mode 100644
index 00000000000..b95046b2b42
--- /dev/null
+++ b/spec/support/helpers/expect_next_instance_of.rb
@@ -0,0 +1,13 @@
+module ExpectNextInstanceOf
+  def expect_next_instance_of(klass, *new_args)
+    receive_new = receive(:new)
+    receive_new.with(*new_args) if new_args.any?
+
+    expect(klass).to receive_new
+      .and_wrap_original do |method, *original_args|
+        method.call(*original_args).tap do |instance|
+          yield(instance)
+        end
+      end
+  end
+end
diff --git a/spec/workers/update_merge_requests_worker_spec.rb b/spec/workers/update_merge_requests_worker_spec.rb
index 80137815d2b..0b553db0ca4 100644
--- a/spec/workers/update_merge_requests_worker_spec.rb
+++ b/spec/workers/update_merge_requests_worker_spec.rb
@@ -18,13 +18,9 @@ describe UpdateMergeRequestsWorker do
     end
 
     it 'executes MergeRequests::RefreshService with expected values' do
-      expect(MergeRequests::RefreshService).to receive(:new)
-        .with(project, user).and_wrap_original do |method, *args|
-          method.call(*args).tap do |refresh_service|
-            expect(refresh_service)
-              .to receive(:execute).with(oldrev, newrev, ref)
-          end
-        end
+      expect_next_instance_of(MergeRequests::RefreshService, project, user) do |refresh_service|
+        expect(refresh_service).to receive(:execute).with(oldrev, newrev, ref)
+      end
 
       perform
     end