1 files changed, 57 insertions, 3 deletions

diff --git a/spec/lib/gitlab/instrumentation_helper_spec.rb b/spec/lib/gitlab/instrumentation_helper_spec.rb
index c00b0fdf043..a5c9cde4c37 100644
--- a/spec/lib/gitlab/instrumentation_helper_spec.rb
+++ b/spec/lib/gitlab/instrumentation_helper_spec.rb
@@ -9,12 +9,17 @@ RSpec.describe Gitlab::InstrumentationHelper do
   describe '.keys' do
     it 'returns all available payload keys' do
       expected_keys = [
+        :cpu_s,
         :gitaly_calls,
         :gitaly_duration_s,
         :rugged_calls,
         :rugged_duration_s,
         :elasticsearch_calls,
         :elasticsearch_duration_s,
+        :elasticsearch_timed_out_count,
+        :mem_objects,
+        :mem_bytes,
+        :mem_mallocs,
         :redis_calls,
         :redis_duration_s,
         :redis_read_bytes,
@@ -37,7 +42,11 @@ RSpec.describe Gitlab::InstrumentationHelper do
         :redis_shared_state_write_bytes,
         :db_count,
         :db_write_count,
-        :db_cached_count
+        :db_cached_count,
+        :external_http_count,
+        :external_http_duration_s,
+        :rack_attack_redis_count,
+        :rack_attack_redis_duration_s
       ]
 
       expect(described_class.keys).to eq(expected_keys)
@@ -49,10 +58,14 @@ RSpec.describe Gitlab::InstrumentationHelper do
 
     subject { described_class.add_instrumentation_data(payload) }
 
-    it 'adds only DB counts by default' do
+    before do
+      described_class.init_instrumentation_data
+    end
+
+    it 'includes DB counts' do
       subject
 
-      expect(payload).to eq(db_count: 0, db_cached_count: 0, db_write_count: 0)
+      expect(payload).to include(db_count: 0, db_cached_count: 0, db_write_count: 0)
     end
 
     context 'when Gitaly calls are made' do
@@ -110,6 +123,47 @@ RSpec.describe Gitlab::InstrumentationHelper do
         expect(payload[:throttle_safelist]).to eq('foobar')
       end
     end
+
+    it 'logs cpu_s duration' do
+      subject
+
+      expect(payload).to include(:cpu_s)
+    end
+
+    context 'when logging memory allocations' do
+      include MemoryInstrumentationHelper
+
+      before do
+        skip_memory_instrumentation!
+      end
+
+      it 'logs memory usage metrics' do
+        subject
+
+        expect(payload).to include(
+          :mem_objects,
+          :mem_bytes,
+          :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
   end
 
   describe '.queue_duration_for_job' do