Fixed stubbing for Gitlab::Metrics specs

If the measure method uses Transaction.current directly the SQL
subscriber (Subscribers::ActiveRecord) will add timings of queries
triggered by DB cleaner.

2 files changed, 13 insertions, 5 deletions

diff --git a/lib/gitlab/metrics.rb b/lib/gitlab/metrics.rb
index 33dd3e39f4d..2a0a5629be5 100644
--- a/lib/gitlab/metrics.rb
+++ b/lib/gitlab/metrics.rb
@@ -82,7 +82,9 @@ module Gitlab
     #
     # Returns the value yielded by the supplied block.
     def self.measure(name)
-      return yield unless Transaction.current
+      trans = current_transaction
+
+      return yield unless trans
 
       real_start = Time.now.to_f
       cpu_start = System.cpu_time
@@ -95,9 +97,9 @@ module Gitlab
       real_time = (real_stop - real_start) * 1000.0
       cpu_time = cpu_stop - cpu_start
 
-      Transaction.current.increment("#{name}_real_time", real_time)
-      Transaction.current.increment("#{name}_cpu_time", cpu_time)
-      Transaction.current.increment("#{name}_call_count", 1)
+      trans.increment("#{name}_real_time", real_time)
+      trans.increment("#{name}_cpu_time", cpu_time)
+      trans.increment("#{name}_call_count", 1)
 
       retval
     end
@@ -113,5 +115,11 @@ module Gitlab
           new(udp: { host: host, port: port })
       end
     end
+
+    private
+
+    def self.current_transaction
+      Transaction.current
+    end
   end
 end
diff --git a/spec/lib/gitlab/metrics_spec.rb b/spec/lib/gitlab/metrics_spec.rb
index a3b68455260..3dee13e27f4 100644
--- a/spec/lib/gitlab/metrics_spec.rb
+++ b/spec/lib/gitlab/metrics_spec.rb
@@ -74,7 +74,7 @@ describe Gitlab::Metrics do
       let(:transaction) { Gitlab::Metrics::Transaction.new }
 
       before do
-        allow(Gitlab::Metrics::Transaction).to receive(:current).
+        allow(Gitlab::Metrics).to receive(:current_transaction).
           and_return(transaction)
       end