1 files changed, 12 insertions, 8 deletions
diff --git a/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb b/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb
index 67336cf83e6..ea9e8fa6795 100644
--- a/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb
+++ b/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb
@@ -10,17 +10,19 @@ describe Gitlab::Metrics::SidekiqMiddleware do
     it 'tracks the transaction' do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
-      expect(Gitlab::Metrics::BackgroundTransaction).to receive(:new)
-        .with(worker.class)
-        .and_call_original
+      expect_next_instance_of(Gitlab::Metrics::BackgroundTransaction) do |transaction|
+        expect(transaction).to receive(:set).with(:sidekiq_queue_duration, instance_of(Float))
+        expect(transaction).to receive(:increment).with(:db_count, 1)
+      end
 
-      expect_any_instance_of(Gitlab::Metrics::Transaction).to receive(:set)
-        .with(:sidekiq_queue_duration, instance_of(Float))
+      middleware.call(worker, message, :test) do
+        ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+      end
 
-      middleware.call(worker, message, :test) { nil }
+      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
     end
 
-    it 'tracks the transaction (for messages without `enqueued_at`)' do
+    it 'tracks the transaction (for messages without `enqueued_at`)', :aggregate_failures do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
       expect(Gitlab::Metrics::BackgroundTransaction).to receive(:new)
@@ -33,7 +35,7 @@ describe Gitlab::Metrics::SidekiqMiddleware do
       middleware.call(worker, {}, :test) { nil }
     end
 
-    it 'tracks any raised exceptions' do
+    it 'tracks any raised exceptions', :aggregate_failures do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
       expect_any_instance_of(Gitlab::Metrics::Transaction)
@@ -44,6 +46,8 @@ describe Gitlab::Metrics::SidekiqMiddleware do
 
       expect { middleware.call(worker, message, :test) }
         .to raise_error(RuntimeError)
+
+      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
     end
   end
 end