1 files changed, 41 insertions, 0 deletions
diff --git a/spec/lib/gitlab/metrics_spec.rb b/spec/lib/gitlab/metrics_spec.rb
index 93dd958168d..8f63a5f2043 100644
--- a/spec/lib/gitlab/metrics_spec.rb
+++ b/spec/lib/gitlab/metrics_spec.rb
@@ -60,4 +60,45 @@ describe Gitlab::Metrics do
       expect(described_class.escape_value(10)).to eq('10')
     end
   end
+
+  describe '.measure' do
+    context 'without a transaction' do
+      it 'returns the return value of the block' do
+        val = Gitlab::Metrics.measure(:foo) { 10 }
+
+        expect(val).to eq(10)
+      end
+    end
+
+    context 'with a transaction' do
+      let(:transaction) { Gitlab::Metrics::Transaction.new }
+
+      before do
+        allow(Gitlab::Metrics::Transaction).to receive(:current).
+          and_return(transaction)
+      end
+
+      it 'adds a metric to the current transaction' do
+        expect(transaction).to receive(:add_metric).
+          with(:foo, { duration: a_kind_of(Numeric) }, { tag: 'value' })
+
+        Gitlab::Metrics.measure(:foo, {}, tag: 'value') { 10 }
+      end
+
+      it 'supports adding of custom values' do
+        values = { duration: a_kind_of(Numeric), number: 10 }
+
+        expect(transaction).to receive(:add_metric).
+          with(:foo, values, { tag: 'value' })
+
+        Gitlab::Metrics.measure(:foo, { number: 10 }, tag: 'value') { 10 }
+      end
+
+      it 'returns the return value of the block' do
+        val = Gitlab::Metrics.measure(:foo) { 10 }
+
+        expect(val).to eq(10)
+      end
+    end
+  end
 end