1 files changed, 31 insertions, 59 deletions

diff --git a/spec/lib/bulk_imports/groups/pipelines/milestones_pipeline_spec.rb b/spec/lib/bulk_imports/groups/pipelines/milestones_pipeline_spec.rb
index f0c34c65257..e5cf75c566b 100644
--- a/spec/lib/bulk_imports/groups/pipelines/milestones_pipeline_spec.rb
+++ b/spec/lib/bulk_imports/groups/pipelines/milestones_pipeline_spec.rb
@@ -5,11 +5,10 @@ require 'spec_helper'
 RSpec.describe BulkImports::Groups::Pipelines::MilestonesPipeline do
   let_it_be(:user) { create(:user) }
   let_it_be(:group) { create(:group) }
-  let_it_be(:cursor) { 'cursor' }
   let_it_be(:timestamp) { Time.new(2020, 01, 01).utc }
   let_it_be(:bulk_import) { create(:bulk_import, user: user) }
 
-  let(:entity) do
+  let_it_be(:entity) do
     create(
       :bulk_import_entity,
       bulk_import: bulk_import,
@@ -20,39 +19,19 @@ RSpec.describe BulkImports::Groups::Pipelines::MilestonesPipeline do
     )
   end
 
-  let(:context) { BulkImports::Pipeline::Context.new(entity) }
+  let_it_be(:tracker) { create(:bulk_import_tracker, entity: entity) }
+  let_it_be(:context) { BulkImports::Pipeline::Context.new(tracker) }
 
   subject { described_class.new(context) }
 
-  def milestone_data(title)
-    {
-      'title' => title,
-      'description' => 'desc',
-      'state' => 'closed',
-      'start_date' => '2020-10-21',
-      'due_date' => '2020-10-22',
-      'created_at' => timestamp.to_s,
-      'updated_at' => timestamp.to_s
-    }
-  end
-
-  def extracted_data(title:, has_next_page:, cursor: nil)
-    page_info = {
-      'end_cursor' => cursor,
-      'has_next_page' => has_next_page
-    }
-
-    BulkImports::Pipeline::ExtractedData.new(data: [milestone_data(title)], page_info: page_info)
-  end
-
   before do
     group.add_owner(user)
   end
 
   describe '#run' do
     it 'imports group milestones' do
-      first_page = extracted_data(title: 'milestone1', has_next_page: true, cursor: cursor)
-      last_page = extracted_data(title: 'milestone2', has_next_page: false)
+      first_page = extracted_data(title: 'milestone1', iid: 1, has_next_page: true)
+      last_page = extracted_data(title: 'milestone2', iid: 2)
 
       allow_next_instance_of(BulkImports::Common::Extractors::GraphqlExtractor) do |extractor|
         allow(extractor)
@@ -75,38 +54,6 @@ RSpec.describe BulkImports::Groups::Pipelines::MilestonesPipeline do
     end
   end
 
-  describe '#after_run' do
-    context 'when extracted data has next page' do
-      it 'updates tracker information and runs pipeline again' do
-        data = extracted_data(title: 'milestone', has_next_page: true, cursor: cursor)
-
-        expect(subject).to receive(:run)
-
-        subject.after_run(data)
-
-        tracker = entity.trackers.find_by(relation: :milestones)
-
-        expect(tracker.has_next_page).to eq(true)
-        expect(tracker.next_page).to eq(cursor)
-      end
-    end
-
-    context 'when extracted data has no next page' do
-      it 'updates tracker information and does not run pipeline' do
-        data = extracted_data(title: 'milestone', has_next_page: false)
-
-        expect(subject).not_to receive(:run)
-
-        subject.after_run(data)
-
-        tracker = entity.trackers.find_by(relation: :milestones)
-
-        expect(tracker.has_next_page).to eq(false)
-        expect(tracker.next_page).to be_nil
-      end
-    end
-  end
-
   describe '#load' do
     it 'creates the milestone' do
       data = milestone_data('milestone')
@@ -120,7 +67,7 @@ RSpec.describe BulkImports::Groups::Pipelines::MilestonesPipeline do
       end
 
       it 'raises NotAllowedError' do
-        data = extracted_data(title: 'milestone', has_next_page: false)
+        data = extracted_data(title: 'milestone')
 
         expect { subject.load(context, data) }.to raise_error(::BulkImports::Pipeline::NotAllowedError)
       end
@@ -148,4 +95,29 @@ RSpec.describe BulkImports::Groups::Pipelines::MilestonesPipeline do
         )
     end
   end
+
+  def milestone_data(title, iid: 1)
+    {
+      'title' => title,
+      'description' => 'desc',
+      'iid' => iid,
+      'state' => 'closed',
+      'start_date' => '2020-10-21',
+      'due_date' => '2020-10-22',
+      'created_at' => timestamp.to_s,
+      'updated_at' => timestamp.to_s
+    }
+  end
+
+  def extracted_data(title:, iid: 1, has_next_page: false)
+    page_info = {
+      'has_next_page' => has_next_page,
+      'next_page' => has_next_page ? 'cursor' : nil
+    }
+
+    BulkImports::Pipeline::ExtractedData.new(
+      data: milestone_data(title, iid: iid),
+      page_info: page_info
+    )
+  end
 end