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

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(EE::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

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(EE::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

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(EE::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

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(EE::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

spec/spec_helper.rb

@@ -71,6 +71,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

spec/support/helpers/expect_next_instance_of.rb

+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

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