1 files changed, 13 insertions, 4 deletions
diff --git a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
index 2fd5dd1d83b..a31686b8061 100644
--- a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
+++ b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
@@ -37,10 +37,11 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
 
         it 'increments only db count value' do
           described_class::DB_COUNTERS.each do |counter|
+            prometheus_counter = "gitlab_transaction_#{counter}_total".to_sym
             if expected_counters[counter] > 0
-              expect(transaction).to receive(:increment).with(counter, 1)
+              expect(transaction).to receive(:increment).with(prometheus_counter, 1)
             else
-              expect(transaction).not_to receive(:increment).with(counter, 1)
+              expect(transaction).not_to receive(:increment).with(prometheus_counter, 1)
             end
           end
 
@@ -74,10 +75,18 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         expect(subscriber).to receive(:current_transaction)
                                 .at_least(:once)
                                 .and_return(transaction)
-        expect(described_class.send(:gitlab_sql_duration_seconds)).to receive(:observe).with({}, 0.002)
+        expect(transaction).to receive(:observe).with(:gitlab_sql_duration_seconds, 0.002)
+
         subscriber.sql(event)
       end
 
+      it 'marks the current thread as using the database' do
+        # since it would already have been toggled by other specs
+        Thread.current[:uses_db_connection] = nil
+
+        expect { subscriber.sql(event) }.to change { Thread.current[:uses_db_connection] }.from(nil).to(true)
+      end
+
       context 'with read query' do
         let(:expected_counters) do
           {
@@ -217,7 +226,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         end
 
         it 'skips schema/begin/commit sql commands' do
-          expect(subscriber).to receive(:current_transaction)
+          allow(subscriber).to receive(:current_transaction)
                                   .at_least(:once)
                                   .and_return(transaction)