1 files changed, 120 insertions, 2 deletions

diff --git a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
index 389b0ef1044..28c3ef229ab 100644
--- a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
+++ b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
@@ -10,6 +10,124 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
   let(:connection) { ActiveRecord::Base.retrieve_connection }
   let(:db_config_name) { ::Gitlab::Database.db_config_name(connection) }
 
+  describe '.load_balancing_metric_counter_keys' do
+    context 'multiple databases' do
+      before do
+        skip_if_multiple_databases_not_setup
+      end
+
+      it 'has expected keys' do
+        expect(described_class.load_balancing_metric_counter_keys).to include(
+          :db_replica_count,
+          :db_primary_count,
+          :db_main_count,
+          :db_main_replica_count,
+          :db_ci_count,
+          :db_ci_replica_count,
+          :db_replica_cached_count,
+          :db_primary_cached_count,
+          :db_main_cached_count,
+          :db_main_replica_cached_count,
+          :db_ci_cached_count,
+          :db_ci_replica_cached_count,
+          :db_replica_wal_count,
+          :db_primary_wal_count,
+          :db_main_wal_count,
+          :db_main_replica_wal_count,
+          :db_ci_wal_count,
+          :db_ci_replica_wal_count,
+          :db_replica_wal_cached_count,
+          :db_primary_wal_cached_count,
+          :db_main_wal_cached_count,
+          :db_main_replica_wal_cached_count,
+          :db_ci_wal_cached_count,
+          :db_ci_replica_wal_cached_count
+        )
+      end
+    end
+
+    context 'single database' do
+      before do
+        skip_if_multiple_databases_are_setup
+      end
+
+      it 'has expected keys' do
+        expect(described_class.load_balancing_metric_counter_keys).to include(
+          :db_replica_count,
+          :db_primary_count,
+          :db_main_count,
+          :db_main_replica_count,
+          :db_replica_cached_count,
+          :db_primary_cached_count,
+          :db_main_cached_count,
+          :db_main_replica_cached_count,
+          :db_replica_wal_count,
+          :db_primary_wal_count,
+          :db_main_wal_count,
+          :db_main_replica_wal_count,
+          :db_replica_wal_cached_count,
+          :db_primary_wal_cached_count,
+          :db_main_wal_cached_count,
+          :db_main_replica_wal_cached_count
+        )
+      end
+
+      it 'does not have ci keys' do
+        expect(described_class.load_balancing_metric_counter_keys).not_to include(
+          :db_ci_count,
+          :db_ci_replica_count,
+          :db_ci_cached_count,
+          :db_ci_replica_cached_count,
+          :db_ci_wal_count,
+          :db_ci_replica_wal_count,
+          :db_ci_wal_cached_count,
+          :db_ci_replica_wal_cached_count
+        )
+      end
+    end
+  end
+
+  describe '.load_balancing_metric_duration_keys' do
+    context 'multiple databases' do
+      before do
+        skip_if_multiple_databases_not_setup
+      end
+
+      it 'has expected keys' do
+        expect(described_class.load_balancing_metric_duration_keys).to include(
+          :db_replica_duration_s,
+          :db_primary_duration_s,
+          :db_main_duration_s,
+          :db_main_replica_duration_s,
+          :db_ci_duration_s,
+          :db_ci_replica_duration_s
+        )
+      end
+    end
+
+    context 'single database' do
+      before do
+        skip_if_multiple_databases_are_setup
+      end
+
+      it 'has expected keys' do
+        expect(described_class.load_balancing_metric_duration_keys).to include(
+          :db_replica_duration_s,
+          :db_primary_duration_s,
+          :db_main_duration_s,
+          :db_main_replica_duration_s
+        )
+      end
+
+      it 'does not have ci keys' do
+        expect(described_class.load_balancing_metric_duration_keys).not_to include(
+          :db_ci_duration_s,
+          :db_ci_replica_duration_s
+        )
+      end
+    end
+  end
+
   describe '#transaction' do
     let(:web_transaction) { double('Gitlab::Metrics::WebTransaction') }
     let(:background_transaction) { double('Gitlab::Metrics::WebTransaction') }
@@ -37,7 +155,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
       end
 
       it 'captures the metrics for web only' do
-        expect(web_transaction).to receive(:observe).with(:gitlab_database_transaction_seconds, 0.23, db_config_name: db_config_name)
+        expect(web_transaction).to receive(:observe).with(:gitlab_database_transaction_seconds, 0.23, { db_config_name: db_config_name })
 
         expect(background_transaction).not_to receive(:observe)
         expect(background_transaction).not_to receive(:increment)
@@ -77,7 +195,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
       end
 
       it 'captures the metrics for web only' do
-        expect(background_transaction).to receive(:observe).with(:gitlab_database_transaction_seconds, 0.23, db_config_name: db_config_name)
+        expect(background_transaction).to receive(:observe).with(:gitlab_database_transaction_seconds, 0.23, { db_config_name: db_config_name })
 
         expect(web_transaction).not_to receive(:observe)
         expect(web_transaction).not_to receive(:increment)