1 files changed, 63 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
index 43e34325419..d881390cd52 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_spec.rb
@@ -19,6 +19,12 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     end
   end
 
+  describe 'validations' do
+    subject { build(:batched_background_migration) }
+
+    it { is_expected.to validate_uniqueness_of(:job_arguments).scoped_to(:job_class_name, :table_name, :column_name) }
+  end
+
   describe '.queue_order' do
     let!(:migration1) { create(:batched_background_migration) }
     let!(:migration2) { create(:batched_background_migration) }
@@ -36,6 +42,38 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
 
     it 'returns the first active migration according to queue order' do
       expect(described_class.active_migration).to eq(migration2)
+      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded)
+    end
+  end
+
+  describe '.queued' do
+    let!(:migration1) { create(:batched_background_migration, :finished) }
+    let!(:migration2) { create(:batched_background_migration, :paused) }
+    let!(:migration3) { create(:batched_background_migration, :active) }
+
+    it 'returns active and paused migrations' do
+      expect(described_class.queued).to contain_exactly(migration2, migration3)
+    end
+  end
+
+  describe '.successful_rows_counts' do
+    let!(:migration1) { create(:batched_background_migration) }
+    let!(:migration2) { create(:batched_background_migration) }
+    let!(:migration_without_jobs) { create(:batched_background_migration) }
+
+    before do
+      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 1000, status: :succeeded)
+      create(:batched_background_migration_job, batched_migration: migration1, batch_size: 200, status: :failed)
+      create(:batched_background_migration_job, batched_migration: migration2, batch_size: 500, status: :succeeded)
+      create(:batched_background_migration_job, batched_migration: migration2, batch_size: 200, status: :running)
+    end
+
+    it 'returns totals from successful jobs' do
+      results = described_class.successful_rows_counts([migration1, migration2, migration_without_jobs])
+
+      expect(results[migration1.id]).to eq(1000)
+      expect(results[migration2.id]).to eq(500)
+      expect(results[migration_without_jobs.id]).to eq(nil)
     end
   end
 
@@ -324,4 +362,29 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
       subject
     end
   end
+
+  describe '.for_configuration' do
+    let!(:migration) do
+      create(
+        :batched_background_migration,
+        job_class_name: 'MyJobClass',
+        table_name: :projects,
+        column_name: :id,
+        job_arguments: [[:id], [:id_convert_to_bigint]]
+      )
+    end
+
+    before do
+      create(:batched_background_migration, job_class_name: 'OtherClass')
+      create(:batched_background_migration, table_name: 'other_table')
+      create(:batched_background_migration, column_name: 'other_column')
+      create(:batched_background_migration, job_arguments: %w[other arguments])
+    end
+
+    it 'finds the migration matching the given configuration parameters' do
+      actual = described_class.for_configuration('MyJobClass', :projects, :id, [[:id], [:id_convert_to_bigint]])
+
+      expect(actual).to contain_exactly(migration)
+    end
+  end
 end