1 files changed, 36 insertions, 0 deletions
diff --git a/spec/lib/gitlab/instrumentation_helper_spec.rb b/spec/lib/gitlab/instrumentation_helper_spec.rb
index 488324ccddc..28ae90d4947 100644
--- a/spec/lib/gitlab/instrumentation_helper_spec.rb
+++ b/spec/lib/gitlab/instrumentation_helper_spec.rb
@@ -117,6 +117,42 @@ RSpec.describe Gitlab::InstrumentationHelper do
         end
       end
     end
+
+    context 'when load balancing is enabled' do
+      include_context 'clear DB Load Balancing configuration'
+
+      before do
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+      end
+
+      it 'includes DB counts' do
+        subject
+
+        expect(payload).to include(db_replica_count: 0,
+                                   db_replica_cached_count: 0,
+                                   db_primary_count: 0,
+                                   db_primary_cached_count: 0,
+                                   db_primary_wal_count: 0,
+                                   db_replica_wal_count: 0)
+      end
+    end
+
+    context 'when load balancing is disabled' do
+      before do
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
+      end
+
+      it 'does not include DB counts' do
+        subject
+
+        expect(payload).not_to include(db_replica_count: 0,
+                                   db_replica_cached_count: 0,
+                                   db_primary_count: 0,
+                                   db_primary_cached_count: 0,
+                                   db_primary_wal_count: 0,
+                                   db_replica_wal_count: 0)
+      end
+    end
   end
 
   describe '.queue_duration_for_job' do