Add latest changes from gitlab-org/gitlab@master

1 files changed, 21 insertions, 8 deletions

diff --git a/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb b/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb
index c800f5c1e06..c66d8b1075c 100644
--- a/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb
+++ b/spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb
@@ -18,8 +18,20 @@ RSpec.describe Gitlab::Metrics::SidekiqMiddleware do
       middleware.call(worker, message, :test) do
         ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
       end
+    end
+
+    it 'prevents database counters from leaking to the next transaction' do
+      worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
-      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
+      2.times do
+        Gitlab::WithRequestStore.with_request_store do
+          middleware.call(worker, message, :test) do
+            ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+          end
+        end
+      end
+
+      expect(message).to include(db_count: 1, db_write_count: 0, db_cached_count: 0)
     end
 
     it 'tracks the transaction (for messages without `enqueued_at`)', :aggregate_failures do
@@ -35,19 +47,20 @@ RSpec.describe Gitlab::Metrics::SidekiqMiddleware do
       middleware.call(worker, {}, :test) { nil }
     end
 
-    it 'tracks any raised exceptions', :aggregate_failures do
+    it 'tracks any raised exceptions', :aggregate_failures, :request_store do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
       expect_any_instance_of(Gitlab::Metrics::Transaction)
-        .to receive(:run).and_raise(RuntimeError)
-
-      expect_any_instance_of(Gitlab::Metrics::Transaction)
         .to receive(:add_event).with(:sidekiq_exception)
 
-      expect { middleware.call(worker, message, :test) }
-        .to raise_error(RuntimeError)
+      expect do
+        middleware.call(worker, message, :test) do
+          ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+          raise RuntimeError
+        end
+      end.to raise_error(RuntimeError)
 
-      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
+      expect(message).to include(db_count: 1, db_write_count: 0, db_cached_count: 0)
     end
   end
 end