1 files changed, 148 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb b/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
index 9f0493ab0d7..779e8e40c97 100644
--- a/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
+++ b/spec/lib/gitlab/database/background_migration/batched_migration_runner_spec.rb
@@ -281,4 +281,152 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigrationRunner do
       end
     end
   end
+
+  describe '#finalize' do
+    let(:migration_wrapper) { Gitlab::Database::BackgroundMigration::BatchedMigrationWrapper.new }
+
+    let(:migration_helpers) { ActiveRecord::Migration.new }
+    let(:table_name) { :_batched_migrations_test_table }
+    let(:column_name) { :some_id }
+    let(:job_arguments) { [:some_id, :some_id_convert_to_bigint] }
+
+    let(:migration_status) { :active }
+
+    let!(:batched_migration) do
+      create(
+        :batched_background_migration,
+        status: migration_status,
+        max_value: 8,
+        batch_size: 2,
+        sub_batch_size: 1,
+        interval: 0,
+        table_name: table_name,
+        column_name: column_name,
+        job_arguments: job_arguments,
+        pause_ms: 0
+      )
+    end
+
+    before do
+      migration_helpers.drop_table table_name, if_exists: true
+      migration_helpers.create_table table_name, id: false do |t|
+        t.integer :some_id, primary_key: true
+        t.integer :some_id_convert_to_bigint
+      end
+
+      migration_helpers.execute("INSERT INTO #{table_name} VALUES (1, 1), (2, 2), (3, NULL), (4, NULL), (5, NULL), (6, NULL), (7, NULL), (8, NULL)")
+    end
+
+    after do
+      migration_helpers.drop_table table_name, if_exists: true
+    end
+
+    context 'when the migration is not yet completed' do
+      before do
+        common_attributes = {
+          batched_migration: batched_migration,
+          batch_size: 2,
+          sub_batch_size: 1,
+          pause_ms: 0
+        }
+
+        create(:batched_background_migration_job, common_attributes.merge(status: :succeeded, min_value: 1, max_value: 2))
+        create(:batched_background_migration_job, common_attributes.merge(status: :pending, min_value: 3, max_value: 4))
+        create(:batched_background_migration_job, common_attributes.merge(status: :failed, min_value: 5, max_value: 6, attempts: 1))
+      end
+
+      it 'completes the migration' do
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
+          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
+          .and_return(batched_migration)
+
+        expect(batched_migration).to receive(:finalizing!).and_call_original
+
+        expect do
+          runner.finalize(
+            batched_migration.job_class_name,
+            table_name,
+            column_name,
+            job_arguments
+          )
+        end.to change { batched_migration.reload.status }.from('active').to('finished')
+
+        expect(batched_migration.batched_jobs).to all(be_succeeded)
+
+        not_converted = migration_helpers.execute("SELECT * FROM #{table_name} WHERE some_id_convert_to_bigint IS NULL")
+        expect(not_converted.to_a).to be_empty
+      end
+
+      context 'when migration fails to complete' do
+        it 'raises an error' do
+          batched_migration.batched_jobs.failed.update_all(attempts: Gitlab::Database::BackgroundMigration::BatchedJob::MAX_ATTEMPTS)
+
+          expect do
+            runner.finalize(
+              batched_migration.job_class_name,
+              table_name,
+              column_name,
+              job_arguments
+            )
+          end.to raise_error described_class::FailedToFinalize
+        end
+      end
+    end
+
+    context 'when the migration is already finished' do
+      let(:migration_status) { :finished }
+
+      it 'is a no-op' do
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
+          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, job_arguments)
+          .and_return(batched_migration)
+
+        configuration = {
+          job_class_name: batched_migration.job_class_name,
+          table_name: table_name.to_sym,
+          column_name: column_name.to_sym,
+          job_arguments: job_arguments
+        }
+
+        expect(Gitlab::AppLogger).to receive(:warn)
+          .with("Batched background migration for the given configuration is already finished: #{configuration}")
+
+        expect(batched_migration).not_to receive(:finalizing!)
+
+        runner.finalize(
+          batched_migration.job_class_name,
+          table_name,
+          column_name,
+          job_arguments
+        )
+      end
+    end
+
+    context 'when the migration does not exist' do
+      it 'is a no-op' do
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_for_configuration)
+          .with('CopyColumnUsingBackgroundMigrationJob', table_name, column_name, [:some, :other, :arguments])
+          .and_return(nil)
+
+        configuration = {
+          job_class_name: batched_migration.job_class_name,
+          table_name: table_name.to_sym,
+          column_name: column_name.to_sym,
+          job_arguments: [:some, :other, :arguments]
+        }
+
+        expect(Gitlab::AppLogger).to receive(:warn)
+          .with("Could not find batched background migration for the given configuration: #{configuration}")
+
+        expect(batched_migration).not_to receive(:finalizing!)
+
+        runner.finalize(
+          batched_migration.job_class_name,
+          table_name,
+          column_name,
+          [:some, :other, :arguments]
+        )
+      end
+    end
+  end
 end