1 files changed, 15 insertions, 14 deletions

diff --git a/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb b/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
index 0282a7af0df..6c32fb3ca17 100644
--- a/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
+++ b/spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb
@@ -5,7 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
   let(:env) { {} }
   let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
-  let(:subject) { described_class.new(env: env, logger: logger, timing_configuration: timing_configuration) }
+  let(:subject) { described_class.new(connection: connection, env: env, logger: logger, timing_configuration: timing_configuration) }
+  let(:connection) { ActiveRecord::Base.retrieve_connection }
 
   let(:timing_configuration) do
     [
@@ -67,7 +68,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
             WHERE t.relkind = 'r' AND l.mode = 'ExclusiveLock' AND t.relname = '#{Project.table_name}'
           """
 
-          expect(ActiveRecord::Base.connection.execute(check_exclusive_lock_query).to_a).to be_present
+          expect(connection.execute(check_exclusive_lock_query).to_a).to be_present
         end
       end
 
@@ -96,8 +97,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
               lock_fiber.resume
             end
 
-            ActiveRecord::Base.transaction do
-              ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+            connection.transaction do
+              connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
               lock_acquired = true
             end
           end
@@ -115,7 +116,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
         context 'setting the idle transaction timeout' do
           context 'when there is no outer transaction: disable_ddl_transaction! is set in the migration' do
             it 'does not disable the idle transaction timeout' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+              allow(connection).to receive(:transaction_open?).and_return(false)
               allow(subject).to receive(:run_block_with_lock_timeout).once.and_raise(ActiveRecord::LockWaitTimeout)
               allow(subject).to receive(:run_block_with_lock_timeout).once
 
@@ -127,7 +128,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
 
           context 'when there is outer transaction: disable_ddl_transaction! is not set in the migration' do
             it 'disables the idle transaction timeout so the code can sleep and retry' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(true)
+              allow(connection).to receive(:transaction_open?).and_return(true)
 
               n = 0
               allow(subject).to receive(:run_block_with_lock_timeout).twice do
@@ -184,8 +185,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
             subject.run(raise_on_exhaustion: true) do
               lock_attempts += 1
 
-              ActiveRecord::Base.transaction do
-                ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+              connection.transaction do
+                connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
                 lock_acquired = true
               end
             end
@@ -199,11 +200,11 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
       context 'when statement timeout is reached' do
         it 'raises StatementInvalid error' do
           lock_acquired = false
-          ActiveRecord::Base.connection.execute("SET statement_timeout='100ms'")
+          connection.execute("SET statement_timeout='100ms'")
 
           expect do
             subject.run do
-              ActiveRecord::Base.connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
+              connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
               lock_acquired = true
             end
           end.to raise_error(ActiveRecord::StatementInvalid)
@@ -216,11 +217,11 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
 
   context 'restore local database variables' do
     it do
-      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW lock_timeout").to_a }
+      expect { subject.run {} }.not_to change { 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 }
+      expect { subject.run {} }.not_to change { connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
     end
   end
 
@@ -228,8 +229,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
     let(:timing_configuration) { [[0.015.seconds, 0.025.seconds], [0.015.seconds, 0.025.seconds]] } # 15ms, 25ms
 
     it 'executes `SET lock_timeout` using the configured timeout value in milliseconds' do
-      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 lock_timeout TO '15ms'").and_call_original
+      expect(connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original
+      expect(connection).to receive(:execute).with("SET lock_timeout TO '15ms'").and_call_original
 
       subject.run { }
     end