1 files changed, 4 insertions, 2 deletions
diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
index 44ef0b307fe..147637cf471 100644
--- a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
@@ -213,7 +213,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
         it 'creates partitions including the next month from today' do
           today = Date.new(2020, 5, 8)
 
-          Timecop.freeze(today) do
+          travel_to(today) do
             migration.partition_table_by_date source_table, partition_column, min_date: min_date
 
             expect_range_partitions_for(partitioned_table, {
@@ -233,7 +233,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
       context 'without min_date, max_date' do
         it 'creates partitions for the current and next month' do
           current_date = Date.new(2020, 05, 22)
-          Timecop.freeze(current_date.to_time) do
+          travel_to(current_date.to_time) do
             migration.partition_table_by_date source_table, partition_column
 
             expect_range_partitions_for(partitioned_table, {
@@ -514,6 +514,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
         allow(migration).to receive(:table_exists?).with(partitioned_table).and_return(true)
         allow(migration).to receive(:copy_missed_records)
         allow(migration).to receive(:execute).with(/VACUUM/)
+        allow(migration).to receive(:execute).with(/^(RE)?SET/)
       end
 
       it 'finishes remaining jobs for the correct table' do
@@ -567,6 +568,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
 
         allow(Gitlab::BackgroundMigration).to receive(:steal)
         allow(migration).to receive(:execute).with(/VACUUM/)
+        allow(migration).to receive(:execute).with(/^(RE)?SET/)
       end
 
       it 'idempotently cleans up after failed background migrations' do