1 files changed, 33 insertions, 38 deletions

diff --git a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
index 54050a87af0..f683ade978a 100644
--- a/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
@@ -58,8 +58,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       it 'does not pass database locations', :aggregate_failures do
         run_middleware
 
-        expect(job['database_replica_location']).to be_nil
-        expect(job['database_write_location']).to be_nil
+        expect(job['wal_locations']).to be_nil
       end
 
       include_examples 'job data consistency'
@@ -86,11 +85,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes database_replica_location' do
+          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
+
           expect(load_balancer).to receive_message_chain(:host, "database_replica_location").and_return(location)
 
           run_middleware
 
-          expect(job['database_replica_location']).to eq(location)
+          expect(job['wal_locations']).to eq(expected_location)
         end
 
         include_examples 'job data consistency'
@@ -102,40 +103,56 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes primary write location', :aggregate_failures do
+          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
+
           expect(load_balancer).to receive(:primary_write_location).and_return(location)
 
           run_middleware
 
-          expect(job['database_write_location']).to eq(location)
+          expect(job['wal_locations']).to eq(expected_location)
         end
 
         include_examples 'job data consistency'
       end
     end
 
-    shared_examples_for 'database location was already provided' do |provided_database_location, other_location|
-      shared_examples_for 'does not set database location again' do |use_primary|
-        before do
-          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(use_primary)
-        end
+    context 'when worker cannot be constantized' do
+      let(:worker_class) { 'ActionMailer::MailDeliveryJob' }
+      let(:expected_consistency) { :always }
 
-        it 'does not set database locations again' do
-          run_middleware
+      include_examples 'does not pass database locations'
+    end
 
-          expect(job[provided_database_location]).to eq(old_location)
-          expect(job[other_location]).to be_nil
-        end
-      end
+    context 'when worker class does not include ApplicationWorker' do
+      let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
+      let(:expected_consistency) { :always }
+
+      include_examples 'does not pass database locations'
+    end
 
+    context 'database wal location was already provided' do
       let(:old_location) { '0/D525E3A8' }
       let(:new_location) { 'AB/12345' }
-      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", provided_database_location => old_location } }
+      let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => old_location } }
+      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations } }
 
       before do
         allow(load_balancer).to receive(:primary_write_location).and_return(new_location)
         allow(load_balancer).to receive(:database_replica_location).and_return(new_location)
       end
 
+      shared_examples_for 'does not set database location again' do |use_primary|
+        before do
+          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(use_primary)
+        end
+
+        it 'does not set database locations again' do
+          run_middleware
+
+          expect(job['wal_locations']).to eq(wal_locations)
+        end
+      end
+
       context "when write was performed" do
         include_examples 'does not set database location again', true
       end
@@ -145,28 +162,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       end
     end
 
-    context 'when worker cannot be constantized' do
-      let(:worker_class) { 'ActionMailer::MailDeliveryJob' }
-      let(:expected_consistency) { :always }
-
-      include_examples 'does not pass database locations'
-    end
-
-    context 'when worker class does not include ApplicationWorker' do
-      let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
-      let(:expected_consistency) { :always }
-
-      include_examples 'does not pass database locations'
-    end
-
-    context 'database write location was already provided' do
-      include_examples 'database location was already provided', 'database_write_location', 'database_replica_location'
-    end
-
-    context 'database replica location was already provided' do
-      include_examples 'database location was already provided', 'database_replica_location', 'database_write_location'
-    end
-
     context 'when worker data consistency is :always' do
       include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker