1 files changed, 8 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/batch_count_spec.rb b/spec/lib/gitlab/database/batch_count_spec.rb
index 7be84b8f980..e7cb53f2dbd 100644
--- a/spec/lib/gitlab/database/batch_count_spec.rb
+++ b/spec/lib/gitlab/database/batch_count_spec.rb
@@ -35,6 +35,10 @@ describe Gitlab::Database::BatchCount do
       expect(described_class.batch_count(model, "#{model.table_name}.id")).to eq(5)
     end
 
+    it 'counts with Arel column' do
+      expect(described_class.batch_count(model, model.arel_table[:id])).to eq(5)
+    end
+
     it 'counts table with batch_size 50K' do
       expect(described_class.batch_count(model, batch_size: 50_000)).to eq(5)
     end
@@ -98,6 +102,10 @@ describe Gitlab::Database::BatchCount do
       expect(described_class.batch_distinct_count(model, "#{model.table_name}.#{column}")).to eq(2)
     end
 
+    it 'counts with Arel column' do
+      expect(described_class.batch_distinct_count(model, model.arel_table[column])).to eq(2)
+    end
+
     it 'counts with :column field with batch_size of 50K' do
       expect(described_class.batch_distinct_count(model, column, batch_size: 50_000)).to eq(2)
     end