1 files changed, 18 insertions, 3 deletions
diff --git a/spec/lib/gitlab/database/with_lock_retries_spec.rb b/spec/lib/gitlab/database/with_lock_retries_spec.rb
index b6321f2eab1..9c8c9749125 100644
--- a/spec/lib/gitlab/database/with_lock_retries_spec.rb
+++ b/spec/lib/gitlab/database/with_lock_retries_spec.rb
@@ -84,7 +84,7 @@ describe Gitlab::Database::WithLockRetries do
           subject.run do
             lock_attempts += 1
 
-            if lock_attempts == retry_count # we reached the last retry iteration, if we kill the thread, the last try (no lock_timeout) will succeed)
+            if lock_attempts == retry_count # we reached the last retry iteration, if we kill the thread, the last try (no lock_timeout) will succeed
               lock_fiber.resume
             end
 
@@ -106,9 +106,13 @@ describe Gitlab::Database::WithLockRetries do
       end
 
       context 'after the retries, without setting lock_timeout' do
-        let(:retry_count) { timing_configuration.size }
+        let(:retry_count) { timing_configuration.size + 1 }
 
-        it_behaves_like 'retriable exclusive lock on `projects`'
+        it_behaves_like 'retriable exclusive lock on `projects`' do
+          before do
+            expect(subject).to receive(:run_block_without_lock_timeout).and_call_original
+          end
+        end
       end
 
       context 'when statement timeout is reached' do
@@ -129,11 +133,22 @@ describe Gitlab::Database::WithLockRetries do
     end
   end
 
+  context 'restore local database variables' do
+    it do
+      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW lock_timeout").to_a }
+    end
+
+    it do
+      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
+    end
+  end
+
   context 'casting durations correctly' do
     let(:timing_configuration) { [[0.015.seconds, 0.025.seconds], [0.015.seconds, 0.025.seconds]] } # 15ms, 25ms
 
     it 'executes `SET LOCAL lock_timeout` using the configured timeout value in milliseconds' do
       expect(ActiveRecord::Base.connection).to receive(:execute).with("SAVEPOINT active_record_1").and_call_original
+      expect(ActiveRecord::Base.connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original
       expect(ActiveRecord::Base.connection).to receive(:execute).with("SET LOCAL lock_timeout TO '15ms'").and_call_original
       expect(ActiveRecord::Base.connection).to receive(:execute).with("RELEASE SAVEPOINT active_record_1").and_call_original