2 files changed, 17 insertions, 5 deletions
diff --git a/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb b/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb
index 0e804b4feac..cd3a94f5737 100644
--- a/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/convert_table_to_first_list_partition_spec.rb
@@ -16,6 +16,7 @@ RSpec.describe Gitlab::Database::Partitioning::ConvertTableToFirstListPartition
   let(:referenced_table_name) { '_test_referenced_table' }
   let(:other_referenced_table_name) { '_test_other_referenced_table' }
   let(:parent_table_name) { "#{table_name}_parent" }
+  let(:lock_tables) { [] }
 
   let(:model) { define_batchable_model(table_name, connection: connection) }
 
@@ -27,7 +28,8 @@ RSpec.describe Gitlab::Database::Partitioning::ConvertTableToFirstListPartition
       table_name: table_name,
       partitioning_column: partitioning_column,
       parent_table_name: parent_table_name,
-      zero_partition_value: partitioning_default
+      zero_partition_value: partitioning_default,
+      lock_tables: lock_tables
     )
   end
 
@@ -168,6 +170,16 @@ RSpec.describe Gitlab::Database::Partitioning::ConvertTableToFirstListPartition
       end
     end
 
+    context 'with locking tables' do
+      let(:lock_tables) { [table_name] }
+
+      it 'locks the table' do
+        recorder = ActiveRecord::QueryRecorder.new { partition }
+
+        expect(recorder.log).to include(/LOCK "_test_table_to_partition" IN ACCESS EXCLUSIVE MODE/)
+      end
+    end
+
     context 'when an error occurs during the conversion' do
       def fail_first_time
         # We can't directly use a boolean here, as we need something that will be passed by-reference to the proc
diff --git a/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb b/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
index 2ef873e8adb..336dec3a8a0 100644
--- a/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb
@@ -92,11 +92,11 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
 
       context 'removing foreign keys' do
         it 'removes foreign keys from the table before dropping it' do
-          expect(dropper).to receive(:drop_detached_partition).and_wrap_original do |drop_method, partition_name|
-            expect(partition_name).to eq('test_partition')
-            expect(foreign_key_exists_by_name(partition_name, 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_falsey
+          expect(dropper).to receive(:drop_detached_partition).and_wrap_original do |drop_method, partition|
+            expect(partition.table_name).to eq('test_partition')
+            expect(foreign_key_exists_by_name(partition.table_name, 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_falsey
 
-            drop_method.call(partition_name)
+            drop_method.call(partition)
           end
 
           expect(foreign_key_exists_by_name('test_partition', 'fk_referenced', schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA)).to be_truthy