1 files changed, 34 insertions, 19 deletions
diff --git a/spec/lib/gitlab/instrumentation_helper_spec.rb b/spec/lib/gitlab/instrumentation_helper_spec.rb
index 28ae90d4947..48fcc9f93db 100644
--- a/spec/lib/gitlab/instrumentation_helper_spec.rb
+++ b/spec/lib/gitlab/instrumentation_helper_spec.rb
@@ -99,23 +99,6 @@ RSpec.describe Gitlab::InstrumentationHelper do
           :mem_mallocs
         )
       end
-
-      context 'when trace_memory_allocations is disabled' do
-        before do
-          stub_feature_flags(trace_memory_allocations: false)
-          Gitlab::Memory::Instrumentation.ensure_feature_flag!
-        end
-
-        it 'does not log memory usage metrics' do
-          subject
-
-          expect(payload).not_to include(
-            :mem_objects,
-            :mem_bytes,
-            :mem_mallocs
-          )
-        end
-      end
     end
 
     context 'when load balancing is enabled' do
@@ -133,7 +116,37 @@ RSpec.describe Gitlab::InstrumentationHelper do
                                    db_primary_count: 0,
                                    db_primary_cached_count: 0,
                                    db_primary_wal_count: 0,
-                                   db_replica_wal_count: 0)
+                                   db_replica_wal_count: 0,
+                                   db_primary_wal_cached_count: 0,
+                                   db_replica_wal_cached_count: 0)
+      end
+
+      context 'when replica caught up search was made' do
+        before do
+          Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 2
+          Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = 1
+        end
+
+        it 'includes related metrics' do
+          subject
+
+          expect(payload).to include(caught_up_replica_pick_ok: 2)
+          expect(payload).to include(caught_up_replica_pick_fail: 1)
+        end
+      end
+
+      context 'when only a single counter was updated' do
+        before do
+          Gitlab::SafeRequestStore[:caught_up_replica_pick_ok] = 1
+          Gitlab::SafeRequestStore[:caught_up_replica_pick_fail] = nil
+        end
+
+        it 'includes only that counter into logging' do
+          subject
+
+          expect(payload).to include(caught_up_replica_pick_ok: 1)
+          expect(payload).not_to include(:caught_up_replica_pick_fail)
+        end
       end
     end
 
@@ -150,7 +163,9 @@ RSpec.describe Gitlab::InstrumentationHelper do
                                    db_primary_count: 0,
                                    db_primary_cached_count: 0,
                                    db_primary_wal_count: 0,
-                                   db_replica_wal_count: 0)
+                                   db_replica_wal_count: 0,
+                                   db_primary_wal_cached_count: 0,
+                                   db_replica_wal_cached_count: 0)
       end
     end
   end