1 files changed, 16 insertions, 2 deletions
diff --git a/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb b/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb
index 076fb9e8215..50ad77caaf1 100644
--- a/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb
+++ b/spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb
@@ -3,7 +3,8 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
   describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigrationProxyLockRetries do
-    let(:migration) { double }
+    let(:connection) { ActiveRecord::Base.connection }
+    let(:migration) { double(connection: connection) }
     let(:return_value) { double }
     let(:class_def) do
       Class.new do
@@ -40,6 +41,18 @@ RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
         expect(result).to eq(return_value)
       end
     end
+
+    describe '#migration_connection' do
+      subject { class_def.new(migration).migration_connection }
+
+      it 'retrieves actual migration connection from #migration' do
+        expect(migration).to receive(:connection).and_return(return_value)
+
+        result = subject
+
+        expect(result).to eq(return_value)
+      end
+    end
   end
 
   describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigratorLockRetries do
@@ -96,7 +109,8 @@ RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
 
     context 'with transactions enabled and lock retries enabled' do
       let(:receiver) { double('receiver', use_transaction?: true)}
-      let(:migration) { double('migration', enable_lock_retries?: true) }
+      let(:migration) { double('migration', migration_connection: connection, enable_lock_retries?: true) }
+      let(:connection) { ActiveRecord::Base.connection }
 
       it 'calls super method' do
         p = proc { }