1 files changed, 27 insertions, 9 deletions

diff --git a/spec/lib/gitlab/database/reindexing/concurrent_reindex_spec.rb b/spec/lib/gitlab/database/reindexing/concurrent_reindex_spec.rb
index 51fc7c6620b..d9077969003 100644
--- a/spec/lib/gitlab/database/reindexing/concurrent_reindex_spec.rb
+++ b/spec/lib/gitlab/database/reindexing/concurrent_reindex_spec.rb
@@ -111,7 +111,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
       end
 
       it 'replaces the existing index with an identical index' do
-        expect(connection).to receive(:execute).with('SET statement_timeout TO \'21600s\'').twice
+        expect(connection).to receive(:execute).with('SET statement_timeout TO \'32400s\'')
 
         expect_to_execute_concurrently_in_order(create_index)
 
@@ -123,6 +123,10 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
         expect_index_rename(replacement_name, index.name)
         expect_index_rename(replaced_name, replacement_name)
 
+        expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+          expect(instance).to receive(:run).with(raise_on_exhaustion: false).and_yield
+        end
+
         expect_to_execute_concurrently_in_order(drop_index)
 
         subject.perform
@@ -136,7 +140,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
         end
 
         it 'rebuilds table statistics before dropping the original index' do
-          expect(connection).to receive(:execute).with('SET statement_timeout TO \'21600s\'').twice
+          expect(connection).to receive(:execute).with('SET statement_timeout TO \'32400s\'')
 
           expect_to_execute_concurrently_in_order(create_index)
 
@@ -152,7 +156,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
           expect_index_rename(replacement_name, index.name)
           expect_index_rename(replaced_name, replacement_name)
 
-          expect_to_execute_concurrently_in_order(drop_index)
+          expect_index_drop(drop_index)
 
           subject.perform
 
@@ -166,9 +170,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
         end
 
         it 'replaces the existing index with an identical index' do
-          expect(connection).to receive(:execute).with('SET statement_timeout TO \'21600s\'').exactly(3).times
-
-          expect_to_execute_concurrently_in_order(drop_index)
+          expect_index_drop(drop_index)
           expect_to_execute_concurrently_in_order(create_index)
 
           expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
@@ -179,7 +181,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
           expect_index_rename(replacement_name, index.name)
           expect_index_rename(replaced_name, replacement_name)
 
-          expect_to_execute_concurrently_in_order(drop_index)
+          expect_index_drop(drop_index)
 
           subject.perform
 
@@ -192,6 +194,10 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
           expect(connection).to receive(:execute).with(create_index).ordered
             .and_raise(ActiveRecord::ConnectionTimeoutError, 'connect timeout')
 
+          expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+            expect(instance).to receive(:run).with(raise_on_exhaustion: false).and_yield
+          end
+
           expect_to_execute_concurrently_in_order(drop_index)
 
           expect { subject.perform }.to raise_error(ActiveRecord::ConnectionTimeoutError, /connect timeout/)
@@ -207,6 +213,10 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
 
           expect_to_execute_concurrently_in_order(create_index)
 
+          expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+            expect(instance).to receive(:run).with(raise_on_exhaustion: false).and_yield
+          end
+
           expect_to_execute_concurrently_in_order(drop_index)
 
           expect { subject.perform }.to raise_error(described_class::ReindexError, /replacement index was created as INVALID/)
@@ -228,7 +238,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
 
           expect_index_rename(index.name, replaced_name).and_raise(ActiveRecord::ConnectionTimeoutError, 'connect timeout')
 
-          expect_to_execute_concurrently_in_order(drop_index)
+          expect_index_drop(drop_index)
 
           expect { subject.perform }.to raise_error(ActiveRecord::ConnectionTimeoutError, /connect timeout/)
 
@@ -245,7 +255,7 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
               .and_raise(::Gitlab::Database::WithLockRetries::AttemptsExhaustedError, 'exhausted')
           end
 
-          expect_to_execute_concurrently_in_order(drop_index)
+          expect_index_drop(drop_index)
 
           expect { subject.perform }.to raise_error(::Gitlab::Database::WithLockRetries::AttemptsExhaustedError, /exhausted/)
 
@@ -270,6 +280,14 @@ RSpec.describe Gitlab::Database::Reindexing::ConcurrentReindex, '#perform' do
     SQL
   end
 
+  def expect_index_drop(drop_index)
+    expect_next_instance_of(::Gitlab::Database::WithLockRetries) do |instance|
+      expect(instance).to receive(:run).with(raise_on_exhaustion: false).and_yield
+    end
+
+    expect_to_execute_concurrently_in_order(drop_index)
+  end
+
   def find_index_create_statement
     ActiveRecord::Base.connection.select_value(<<~SQL)
       SELECT indexdef