1 files changed, 40 insertions, 20 deletions
diff --git a/spec/workers/gitlab/github_import/stage/import_attachments_worker_spec.rb b/spec/workers/gitlab/github_import/stage/import_attachments_worker_spec.rb
index c2c5e1dbf4e..ecfece735af 100644
--- a/spec/workers/gitlab/github_import/stage/import_attachments_worker_spec.rb
+++ b/spec/workers/gitlab/github_import/stage/import_attachments_worker_spec.rb
@@ -5,40 +5,60 @@ require 'spec_helper'
 RSpec.describe Gitlab::GithubImport::Stage::ImportAttachmentsWorker do
   subject(:worker) { described_class.new }
 
-  let(:project) { create(:project) }
-  let!(:group) { create(:group, projects: [project]) }
-  let(:feature_flag_state) { [group] }
+  let_it_be(:project) { create(:project) }
+  let(:settings) { ::Gitlab::GithubImport::Settings.new(project) }
+  let(:stage_enabled) { true }
+
+  before do
+    settings.write({ attachments_import: stage_enabled })
+  end
 
   describe '#import' do
-    let(:importer) { instance_double('Gitlab::GithubImport::Importer::ReleasesAttachmentsImporter') }
     let(:client) { instance_double('Gitlab::GithubImport::Client') }
-
-    before do
-      stub_feature_flags(github_importer_attachments_import: feature_flag_state)
+    let(:importers) do
+      [
+        {
+          klass: Gitlab::GithubImport::Importer::Attachments::ReleasesImporter,
+          double: instance_double('Gitlab::GithubImport::Importer::Attachments::ReleasesImporter'),
+          waiter: Gitlab::JobWaiter.new(2, '123')
+        },
+        {
+          klass: Gitlab::GithubImport::Importer::Attachments::NotesImporter,
+          double: instance_double('Gitlab::GithubImport::Importer::Attachments::NotesImporter'),
+          waiter: Gitlab::JobWaiter.new(3, '234')
+        },
+        {
+          klass: Gitlab::GithubImport::Importer::Attachments::IssuesImporter,
+          double: instance_double('Gitlab::GithubImport::Importer::Attachments::IssuesImporter'),
+          waiter: Gitlab::JobWaiter.new(4, '345')
+        },
+        {
+          klass: Gitlab::GithubImport::Importer::Attachments::MergeRequestsImporter,
+          double: instance_double('Gitlab::GithubImport::Importer::Attachments::MergeRequestsImporter'),
+          waiter: Gitlab::JobWaiter.new(5, '456')
+        }
+      ]
     end
 
-    it 'imports release attachments' do
-      waiter = Gitlab::JobWaiter.new(2, '123')
-
-      expect(Gitlab::GithubImport::Importer::ReleasesAttachmentsImporter)
-        .to receive(:new)
-        .with(project, client)
-        .and_return(importer)
-
-      expect(importer).to receive(:execute).and_return(waiter)
+    it 'imports attachments' do
+      importers.each do |importer|
+        expect_next_instance_of(importer[:klass], project, client) do |instance|
+          expect(instance).to receive(:execute).and_return(importer[:waiter])
+        end
+      end
 
       expect(Gitlab::GithubImport::AdvanceStageWorker)
         .to receive(:perform_async)
-        .with(project.id, { '123' => 2 }, :protected_branches)
+        .with(project.id, { '123' => 2, '234' => 3, '345' => 4, '456' => 5 }, :protected_branches)
 
       worker.import(client, project)
     end
 
-    context 'when feature flag is disabled' do
-      let(:feature_flag_state) { false }
+    context 'when stage is disabled' do
+      let(:stage_enabled) { false }
 
       it 'skips release attachments import and calls next stage' do
-        expect(Gitlab::GithubImport::Importer::ReleasesAttachmentsImporter).not_to receive(:new)
+        importers.each { |importer| expect(importer[:klass]).not_to receive(:new) }
         expect(Gitlab::GithubImport::AdvanceStageWorker)
           .to receive(:perform_async).with(project.id, {}, :protected_branches)