1 files changed, 30 insertions, 0 deletions
diff --git a/spec/lib/gitlab/utils/usage_data_spec.rb b/spec/lib/gitlab/utils/usage_data_spec.rb
index 9c0dc69ccd1..521d6584a20 100644
--- a/spec/lib/gitlab/utils/usage_data_spec.rb
+++ b/spec/lib/gitlab/utils/usage_data_spec.rb
@@ -37,6 +37,36 @@ RSpec.describe Gitlab::Utils::UsageData do
     end
   end
 
+  describe '#estimate_batch_distinct_count' do
+    let(:relation) { double(:relation) }
+
+    it 'delegates counting to counter class instance' do
+      expect_next_instance_of(Gitlab::Database::PostgresHll::BatchDistinctCounter, relation, 'column') do |instance|
+        expect(instance).to receive(:estimate_distinct_count)
+                              .with(batch_size: nil, start: nil, finish: nil)
+                              .and_return(5)
+      end
+
+      expect(described_class.estimate_batch_distinct_count(relation, 'column')).to eq(5)
+    end
+
+    it 'returns default fallback value when counting fails due to database error' do
+      stub_const("Gitlab::Utils::UsageData::FALLBACK", 15)
+      allow(Gitlab::Database::PostgresHll::BatchDistinctCounter).to receive(:new).and_raise(ActiveRecord::StatementInvalid.new(''))
+
+      expect(described_class.estimate_batch_distinct_count(relation)).to eq(15)
+    end
+
+    it 'logs error and returns DISTRIBUTED_HLL_FALLBACK value when counting raises any error', :aggregate_failures do
+      error = StandardError.new('')
+      stub_const("Gitlab::Utils::UsageData::DISTRIBUTED_HLL_FALLBACK", 15)
+      allow(Gitlab::Database::PostgresHll::BatchDistinctCounter).to receive(:new).and_raise(error)
+
+      expect(Gitlab::ErrorTracking).to receive(:track_and_raise_for_dev_exception).with(error)
+      expect(described_class.estimate_batch_distinct_count(relation)).to eq(15)
+    end
+  end
+
   describe '#sum' do
     let(:relation) { double(:relation) }