1 files changed, 34 insertions, 23 deletions
diff --git a/spec/lib/gitlab/lograge/custom_options_spec.rb b/spec/lib/gitlab/lograge/custom_options_spec.rb
index 7ae8baa31b5..ebf150d21ef 100644
--- a/spec/lib/gitlab/lograge/custom_options_spec.rb
+++ b/spec/lib/gitlab/lograge/custom_options_spec.rb
@@ -13,21 +13,16 @@ describe Gitlab::Lograge::CustomOptions do
       }
     end
 
-    let(:event) do
-      ActiveSupport::Notifications::Event.new(
-        'test',
-        1,
-        2,
-        'transaction_id',
-        {
-          params: params,
-          user_id: 'test',
-          cf_ray: SecureRandom.hex,
-          cf_request_id: SecureRandom.hex,
-          metadata: { 'meta.user' => 'jane.doe' }
-        }
-      )
+    let(:event_payload) do
+      {
+        params: params,
+        user_id: 'test',
+        cf_ray: SecureRandom.hex,
+        cf_request_id: SecureRandom.hex,
+        metadata: { 'meta.user' => 'jane.doe' }
+      }
     end
+    let(:event) { ActiveSupport::Notifications::Event.new('test', 1, 2, 'transaction_id', event_payload) }
 
     subject { described_class.call(event) }
 
@@ -49,6 +44,18 @@ describe Gitlab::Lograge::CustomOptions do
       end
     end
 
+    context 'with transaction' do
+      let(:transaction) { Gitlab::Metrics::WebTransaction.new({}) }
+
+      before do
+        allow(Gitlab::Metrics::Transaction).to receive(:current).and_return(transaction)
+      end
+
+      it 'adds db counters' do
+        expect(subject).to include(:db_count, :db_write_count, :db_cached_count)
+      end
+    end
+
     it 'adds the user id' do
       expect(subject[:user_id]).to eq('test')
     end
@@ -63,19 +70,23 @@ describe Gitlab::Lograge::CustomOptions do
     end
 
     context 'when metadata is missing' do
-      let(:event) do
-        ActiveSupport::Notifications::Event.new(
-          'test',
-          1,
-          2,
-          'transaction_id',
-          { params: {} }
-        )
-      end
+      let(:event_payload) { { params: {} } }
 
       it 'does not break' do
         expect { subject }.not_to raise_error
       end
     end
+
+    context 'when correlation_id is overriden' do
+      let(:correlation_id_key) { Labkit::Correlation::CorrelationId::LOG_KEY }
+
+      before do
+        event_payload[correlation_id_key] = '123456'
+      end
+
+      it 'sets the overriden value' do
+        expect(subject[correlation_id_key]).to eq('123456')
+      end
+    end
   end
 end