Schedule stage_id background migration in ranges

1 files changed, 5 insertions, 6 deletions

diff --git a/spec/migrations/migrate_stage_id_reference_in_background_spec.rb b/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
index a32a7fceb68..f3dde8b59c0 100644
--- a/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
+++ b/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
@@ -21,7 +21,7 @@ describe MigrateStageIdReferenceInBackground, :migration, :sidekiq do
   let(:projects) { table(:projects) }
 
   before do
-    stub_const('MigrateStageIdReferenceInBackground::BATCH_SIZE', 2)
+    stub_const("#{described_class.name}::BATCH_SIZE", 2)
 
     projects.create!(id: 123, name: 'gitlab1', path: 'gitlab1')
     projects.create!(id: 345, name: 'gitlab2', path: 'gitlab2')
@@ -47,11 +47,10 @@ describe MigrateStageIdReferenceInBackground, :migration, :sidekiq do
       Timecop.freeze do
         migrate!
 
-        expect(described_class::MIGRATION).to be_scheduled_migration(2.minutes, 1)
-        expect(described_class::MIGRATION).to be_scheduled_migration(2.minutes, 2)
-        expect(described_class::MIGRATION).to be_scheduled_migration(4.minutes, 3)
-        expect(described_class::MIGRATION).to be_scheduled_migration(4.minutes, 4)
-        expect(BackgroundMigrationWorker.jobs.size).to eq 5
+        expect(described_class::MIGRATION).to be_scheduled_migration(2.minutes, 1, 3)
+        expect(described_class::MIGRATION).to be_scheduled_migration(4.minutes, 3, 5)
+        expect(described_class::MIGRATION).to be_scheduled_migration(6.minutes, 5, 0)
+        expect(BackgroundMigrationWorker.jobs.size).to eq 3
       end
     end
   end