1 files changed, 102 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 7e3de32b965..d71a4f81901 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -14,6 +14,54 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     allow(model).to receive(:puts)
   end
 
+  describe 'overridden dynamic model helpers' do
+    let(:test_table) { '__test_batching_table' }
+
+    before do
+      model.connection.execute(<<~SQL)
+        CREATE TABLE #{test_table} (
+          id integer NOT NULL PRIMARY KEY,
+          name text NOT NULL
+        );
+
+        INSERT INTO #{test_table} (id, name)
+        VALUES (1, 'bob'), (2, 'mary'), (3, 'amy');
+      SQL
+    end
+
+    describe '#define_batchable_model' do
+      it 'defines a batchable model with the migration connection' do
+        expect(model.define_batchable_model(test_table).count).to eq(3)
+      end
+    end
+
+    describe '#each_batch' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'calls each_batch with the migration connection' do
+        each_batch_name = ->(&block) do
+          model.each_batch(test_table, of: 2) do |batch|
+            block.call(batch.pluck(:name))
+          end
+        end
+
+        expect { |b| each_batch_name.call(&b) }.to yield_successive_args(%w[bob mary], %w[amy])
+      end
+    end
+
+    describe '#each_batch_range' do
+      before do
+        allow(model).to receive(:transaction_open?).and_return(false)
+      end
+
+      it 'calls each_batch with the migration connection' do
+        expect { |b| model.each_batch_range(test_table, of: 2, &b) }.to yield_successive_args([1, 2], [3, 3])
+      end
+    end
+  end
+
   describe '#remove_timestamps' do
     it 'can remove the default timestamps' do
       Gitlab::Database::MigrationHelpers::DEFAULT_TIMESTAMP_COLUMNS.each do |column_name|
@@ -442,6 +490,60 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
     end
   end
 
+  describe '#remove_foreign_key_if_exists' do
+    context 'when the foreign key does not exist' do
+      before do
+        allow(model).to receive(:foreign_key_exists?).and_return(false)
+      end
+
+      it 'does nothing' do
+        expect(model).not_to receive(:remove_foreign_key)
+
+        model.remove_foreign_key_if_exists(:projects, :users, column: :user_id)
+      end
+    end
+
+    context 'when the foreign key exists' do
+      before do
+        allow(model).to receive(:foreign_key_exists?).and_return(true)
+      end
+
+      it 'removes the foreign key' do
+        expect(model).to receive(:remove_foreign_key).with(:projects, :users, { column: :user_id })
+
+        model.remove_foreign_key_if_exists(:projects, :users, column: :user_id)
+      end
+
+      context 'when the target table is not given' do
+        it 'passes the options as the second parameter' do
+          expect(model).to receive(:remove_foreign_key).with(:projects, { column: :user_id })
+
+          model.remove_foreign_key_if_exists(:projects, column: :user_id)
+        end
+      end
+
+      context 'when the reverse_lock_order option is given' do
+        it 'requests for lock before removing the foreign key' do
+          expect(model).to receive(:transaction_open?).and_return(true)
+          expect(model).to receive(:execute).with(/LOCK TABLE users, projects/)
+          expect(model).not_to receive(:remove_foreign_key).with(:projects, :users)
+
+          model.remove_foreign_key_if_exists(:projects, :users, column: :user_id, reverse_lock_order: true)
+        end
+
+        context 'when not inside a transaction' do
+          it 'does not lock' do
+            expect(model).to receive(:transaction_open?).and_return(false)
+            expect(model).not_to receive(:execute).with(/LOCK TABLE users, projects/)
+            expect(model).to receive(:remove_foreign_key).with(:projects, :users, { column: :user_id })
+
+            model.remove_foreign_key_if_exists(:projects, :users, column: :user_id, reverse_lock_order: true)
+          end
+        end
+      end
+    end
+  end
+
   describe '#add_concurrent_foreign_key' do
     before do
       allow(model).to receive(:foreign_key_exists?).and_return(false)