1 files changed, 54 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/batch_count_spec.rb b/spec/lib/gitlab/database/batch_count_spec.rb
index 028bdce852e..811d4fad95c 100644
--- a/spec/lib/gitlab/database/batch_count_spec.rb
+++ b/spec/lib/gitlab/database/batch_count_spec.rb
@@ -384,4 +384,58 @@ RSpec.describe Gitlab::Database::BatchCount do
       subject { described_class.method(:batch_sum) }
     end
   end
+
+  describe '#batch_average' do
+    let(:model) { Issue }
+    let(:column) { :weight }
+
+    before do
+      Issue.update_all(weight: 2)
+    end
+
+    it 'returns the average of values in the given column' do
+      expect(described_class.batch_average(model, column)).to eq(2)
+    end
+
+    it 'works when given an Arel column' do
+      expect(described_class.batch_average(model, model.arel_table[column])).to eq(2)
+    end
+
+    it 'works with a batch size of 50K' do
+      expect(described_class.batch_average(model, column, batch_size: 50_000)).to eq(2)
+    end
+
+    it 'works with start and finish provided' do
+      expect(described_class.batch_average(model, column, start: model.minimum(:id), finish: model.maximum(:id))).to eq(2)
+    end
+
+    it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE}" do
+      min_id = model.minimum(:id)
+      relation = instance_double(ActiveRecord::Relation)
+      allow(model).to receive_message_chain(:select, public_send: relation)
+      batch_end_id = min_id + calculate_batch_size(Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE)
+
+      expect(relation).to receive(:where).with("id" => min_id..batch_end_id).and_return(double(send: 1))
+
+      described_class.batch_average(model, column)
+    end
+
+    it_behaves_like 'when a transaction is open' do
+      subject { described_class.batch_average(model, column) }
+    end
+
+    it_behaves_like 'disallowed configurations', :batch_average do
+      let(:args) { [model, column] }
+      let(:default_batch_size) { Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE }
+      let(:small_batch_size)   { Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE - 1 }
+    end
+
+    it_behaves_like 'when batch fetch query is canceled' do
+      let(:mode) { :itself }
+      let(:operation) { :average }
+      let(:operation_args) { [column] }
+
+      subject { described_class.method(:batch_average) }
+    end
+  end
 end