1 files changed, 21 insertions, 36 deletions

diff --git a/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb b/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
index 4406b34e638..a98038cd3f8 100644
--- a/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
+++ b/spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb
@@ -228,8 +228,6 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
     end
 
     context 'when the job performs database queries' do
-      include_context 'clear DB Load Balancing configuration'
-
       before do
         allow(Time).to receive(:now).and_return(timestamp)
         allow(Process).to receive(:clock_gettime).and_call_original
@@ -256,7 +254,7 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
           expect(logger).to receive(:info).with(expected_end_payload_with_db).ordered
 
           call_subject(job, 'test_queue') do
-            ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+            ApplicationRecord.connection.execute('SELECT pg_sleep(0.1);')
           end
         end
 
@@ -267,7 +265,7 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
           expect(logger).to receive(:info).with(expected_end_payload).ordered
 
           call_subject(job.dup, 'test_queue') do
-            ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+            ApplicationRecord.connection.execute('SELECT pg_sleep(0.1);')
           end
 
           Gitlab::SafeRequestStore.clear!
@@ -293,54 +291,41 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         include_examples 'performs database queries'
       end
 
-      context 'when load balancing is enabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        end
+      context 'when load balancing is enabled', :db_load_balancing do
+        let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }
 
-        let(:dbname) { ::Gitlab::Database.dbname(ActiveRecord::Base.connection) }
+        let(:expected_db_payload_defaults) do
+          metrics =
+            ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_counter_keys +
+            ::Gitlab::Metrics::Subscribers::ActiveRecord.load_balancing_metric_duration_keys +
+            ::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_keys +
+            [:db_duration_s]
+
+          metrics.each_with_object({}) do |key, result|
+            result[key.to_s] = 0
+          end
+        end
 
         let(:expected_end_payload_with_db) do
-          expected_end_payload.merge(
+          expected_end_payload.merge(expected_db_payload_defaults).merge(
             'db_duration_s' => a_value >= 0.1,
             'db_count' => a_value >= 1,
-            'db_cached_count' => 0,
-            'db_write_count' => 0,
-            'db_replica_count' => 0,
-            'db_replica_cached_count' => 0,
-            'db_replica_wal_count' => 0,
+            "db_replica_#{db_config_name}_count" => 0,
             'db_replica_duration_s' => a_value >= 0,
             'db_primary_count' => a_value >= 1,
-            'db_primary_cached_count' => 0,
-            'db_primary_wal_count' => 0,
+            "db_primary_#{db_config_name}_count" => a_value >= 1,
             'db_primary_duration_s' => a_value > 0,
-            "db_primary_#{dbname}_duration_s" => a_value > 0,
-            'db_primary_wal_cached_count' => 0,
-            'db_replica_wal_cached_count' => 0
+            "db_primary_#{db_config_name}_duration_s" => a_value > 0
           )
         end
 
         let(:end_payload) do
-          start_payload.merge(
+          start_payload.merge(expected_db_payload_defaults).merge(
             'message' => 'TestWorker JID-da883554ee4fe414012f5f42: done: 0.0 sec',
             'job_status' => 'done',
             'duration_s' => 0.0,
             'completed_at' => timestamp.to_f,
-            'cpu_s' => 1.111112,
-            'db_duration_s' => 0.0,
-            'db_cached_count' => 0,
-            'db_count' => 0,
-            'db_write_count' => 0,
-            'db_replica_count' => 0,
-            'db_replica_cached_count' => 0,
-            'db_replica_wal_count' => 0,
-            'db_replica_duration_s' => 0,
-            'db_primary_count' => 0,
-            'db_primary_cached_count' => 0,
-            'db_primary_wal_count' => 0,
-            'db_primary_wal_cached_count' => 0,
-            'db_replica_wal_cached_count' => 0,
-            'db_primary_duration_s' => 0
+            'cpu_s' => 1.111112
           )
         end