Use replace table helper in partitioning tests

Update the partitioning migration tests to also verify that the oid of
the original and replacement tables update as expected, which guarantees
that the tables swapped correctly.

spec/lib/gitlab/database/partitioning/replace_table_spec.rb

@@ -65,14 +65,7 @@ RSpec.describe Gitlab::Database::Partitioning::ReplaceTable, '#perform' do
     expect(primary_key_constraint_name(archived_table)).to eq(archived_primary_key)
   end
 
-  def expect_table_to_be_replaced
-    original_id = table_oid(original_table)
-    replacement_id = table_oid(replacement_table)
-
-    yield
-
-    expect(table_oid(original_table)).to eq(replacement_id)
-    expect(table_oid(archived_table)).to eq(original_id)
-    expect(table_oid(replacement_table)).to be_nil
+  def expect_table_to_be_replaced(&block)
+    super(original_table: original_table, replacement_table: replacement_table, archived_table: archived_table, &block)
   end
 end

spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb

@@ -642,7 +642,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       expect(table_type(partitioned_table)).to eq('partitioned')
       expect(table_type(archived_table)).to be_nil
 
-      migration.replace_with_partitioned_table source_table
+      expect_table_to_be_replaced { migration.replace_with_partitioned_table(source_table) }
 
       expect(table_type(source_table)).to eq('partitioned')
       expect(table_type(archived_table)).to eq('normal')
@@ -653,11 +653,15 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       expect_function_to_exist(function_name)
       expect_valid_function_trigger(source_table, trigger_name, function_name, after: %w[delete insert update])
 
-      migration.replace_with_partitioned_table source_table
+      expect_table_to_be_replaced { migration.replace_with_partitioned_table(source_table) }
 
       expect_function_to_exist(function_name)
       expect_valid_function_trigger(source_table, trigger_name, function_name, after: %w[delete insert update])
     end
+
+    def expect_table_to_be_replaced(&block)
+      super(original_table: source_table, replacement_table: partitioned_table, archived_table: archived_table, &block)
+    end
   end
 
   describe '#rollback_replace_with_partitioned_table' do
@@ -674,7 +678,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       expect(table_type(archived_table)).to eq('normal')
       expect(table_type(partitioned_table)).to be_nil
 
-      migration.rollback_replace_with_partitioned_table source_table
+      expect_table_to_be_replaced { migration.rollback_replace_with_partitioned_table(source_table) }
 
       expect(table_type(source_table)).to eq('normal')
       expect(table_type(partitioned_table)).to eq('partitioned')
@@ -685,11 +689,15 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       expect_function_to_exist(function_name)
       expect_valid_function_trigger(source_table, trigger_name, function_name, after: %w[delete insert update])
 
-      migration.rollback_replace_with_partitioned_table source_table
+      expect_table_to_be_replaced { migration.rollback_replace_with_partitioned_table(source_table) }
 
       expect_function_to_exist(function_name)
       expect_valid_function_trigger(source_table, trigger_name, function_name, after: %w[delete insert update])
     end
+
+    def expect_table_to_be_replaced(&block)
+      super(original_table: source_table, replacement_table: archived_table, archived_table: partitioned_table, &block)
+    end
   end
 
   def filter_columns_by_name(columns, names)

spec/support/helpers/table_schema_helpers.rb

@@ -5,6 +5,17 @@ module TableSchemaHelpers
     ActiveRecord::Base.connection
   end
 
+  def expect_table_to_be_replaced(original_table:, replacement_table:, archived_table:)
+    original_oid = table_oid(original_table)
+    replacement_oid = table_oid(replacement_table)
+
+    yield
+
+    expect(table_oid(original_table)).to eq(replacement_oid)
+    expect(table_oid(archived_table)).to eq(original_oid)
+    expect(table_oid(replacement_table)).to be_nil
+  end
+
   def table_oid(name)
     connection.select_value(<<~SQL)
       SELECT oid