1 files changed, 54 insertions, 14 deletions

diff --git a/spec/lib/gitlab/usage/metrics/instrumentations/database_metric_spec.rb b/spec/lib/gitlab/usage/metrics/instrumentations/database_metric_spec.rb
index 5e36820df5e..0a32bdb95d3 100644
--- a/spec/lib/gitlab/usage/metrics/instrumentations/database_metric_spec.rb
+++ b/spec/lib/gitlab/usage/metrics/instrumentations/database_metric_spec.rb
@@ -4,11 +4,11 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
   subject do
-    described_class.tap do |m|
-      m.relation { Issue }
-      m.operation :count
-      m.start { m.relation.minimum(:id) }
-      m.finish { m.relation.maximum(:id) }
+    described_class.tap do |metric_class|
+      metric_class.relation { Issue }
+      metric_class.operation :count
+      metric_class.start { metric_class.relation.minimum(:id) }
+      metric_class.finish { metric_class.relation.maximum(:id) }
     end.new(time_frame: 'all')
   end
 
@@ -38,9 +38,9 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
 
     context 'with start and finish not called' do
       subject do
-        described_class.tap do |m|
-          m.relation { Issue }
-          m.operation :count
+        described_class.tap do |metric_class|
+          metric_class.relation { Issue }
+          metric_class.operation :count
         end.new(time_frame: 'all')
       end
 
@@ -51,12 +51,12 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
 
     context 'with cache_start_and_finish_as called' do
       subject do
-        described_class.tap do |m|
-          m.relation { Issue }
-          m.operation :count
-          m.start { m.relation.minimum(:id) }
-          m.finish { m.relation.maximum(:id) }
-          m.cache_start_and_finish_as :special_issue_count
+        described_class.tap do |metric_class|
+          metric_class.relation { Issue }
+          metric_class.operation :count
+          metric_class.start { metric_class.relation.minimum(:id) }
+          metric_class.finish { metric_class.relation.maximum(:id) }
+          metric_class.cache_start_and_finish_as :special_issue_count
         end.new(time_frame: 'all')
       end
 
@@ -71,5 +71,45 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
         expect(Rails.cache.read('metric_instrumentation/special_issue_count_maximum_id')).to eq(issues.max_by(&:id).id)
       end
     end
+
+    context 'with estimate_batch_distinct_count' do
+      subject do
+        described_class.tap do |metric_class|
+          metric_class.relation { Issue }
+          metric_class.operation(:estimate_batch_distinct_count)
+          metric_class.start { metric_class.relation.minimum(:id) }
+          metric_class.finish { metric_class.relation.maximum(:id) }
+        end.new(time_frame: 'all')
+      end
+
+      it 'calculates a correct result' do
+        expect(subject.value).to be_within(Gitlab::Database::PostgresHll::BatchDistinctCounter::ERROR_RATE).percent_of(3)
+      end
+
+      context 'with block passed to operation' do
+        let(:buckets) { double('Buckets').as_null_object }
+
+        subject do
+          described_class.tap do |metric_class|
+            metric_class.relation { Issue }
+            metric_class.operation(:estimate_batch_distinct_count) do |result|
+              result.foo
+            end
+            metric_class.start { metric_class.relation.minimum(:id) }
+            metric_class.finish { metric_class.relation.maximum(:id) }
+          end.new(time_frame: 'all')
+        end
+
+        before do
+          allow(Gitlab::Database::PostgresHll::Buckets).to receive(:new).and_return(buckets)
+        end
+
+        it 'calls the block passing HLL buckets as an argument' do
+          expect(buckets).to receive(:foo)
+
+          subject.value
+        end
+      end
+    end
   end
 end