1 files changed, 92 insertions, 7 deletions

diff --git a/spec/models/concerns/each_batch_spec.rb b/spec/models/concerns/each_batch_spec.rb
index f1fb4fcbd03..2c75d4d5c41 100644
--- a/spec/models/concerns/each_batch_spec.rb
+++ b/spec/models/concerns/each_batch_spec.rb
@@ -3,17 +3,17 @@
 require 'spec_helper'
 
 RSpec.describe EachBatch do
-  describe '.each_batch' do
-    let(:model) do
-      Class.new(ActiveRecord::Base) do
-        include EachBatch
+  let(:model) do
+    Class.new(ActiveRecord::Base) do
+      include EachBatch
 
-        self.table_name = 'users'
+      self.table_name = 'users'
 
-        scope :never_signed_in, -> { where(sign_in_count: 0) }
-      end
+      scope :never_signed_in, -> { where(sign_in_count: 0) }
     end
+  end
 
+  describe '.each_batch' do
     before do
       create_list(:user, 5, updated_at: 1.day.ago)
     end
@@ -86,4 +86,89 @@ RSpec.describe EachBatch do
       end
     end
   end
+
+  describe '.distinct_each_batch' do
+    let_it_be(:users) { create_list(:user, 5, sign_in_count: 0) }
+
+    let(:params) { {} }
+
+    subject(:values) do
+      values = []
+
+      model.distinct_each_batch(**params) { |rel| values.concat(rel.pluck(params[:column])) }
+      values
+    end
+
+    context 'when iterating over a unique column' do
+      context 'when using ascending order' do
+        let(:expected_values) { users.pluck(:id).sort }
+        let(:params) { { column: :id, of: 1, order: :asc } }
+
+        it { is_expected.to eq(expected_values) }
+
+        context 'when using larger batch size' do
+          before do
+            params[:of] = 3
+          end
+
+          it { is_expected.to eq(expected_values) }
+        end
+
+        context 'when using larger batch size than the result size' do
+          before do
+            params[:of] = 100
+          end
+
+          it { is_expected.to eq(expected_values) }
+        end
+      end
+
+      context 'when using descending order' do
+        let(:expected_values) { users.pluck(:id).sort.reverse }
+        let(:params) { { column: :id, of: 1, order: :desc } }
+
+        it { is_expected.to eq(expected_values) }
+
+        context 'when using larger batch size' do
+          before do
+            params[:of] = 3
+          end
+
+          it { is_expected.to eq(expected_values) }
+        end
+      end
+    end
+
+    context 'when iterating over a non-unique column' do
+      let(:params) { { column: :sign_in_count, of: 2, order: :asc } }
+
+      context 'when only one value is present' do
+        it { is_expected.to eq([0]) }
+      end
+
+      context 'when duplicated values present' do
+        let(:expected_values) { [2, 5] }
+
+        before do
+          users[0].reload.update!(sign_in_count: 5)
+          users[1].reload.update!(sign_in_count: 2)
+          users[2].reload.update!(sign_in_count: 5)
+          users[3].reload.update!(sign_in_count: 2)
+          users[4].reload.update!(sign_in_count: 5)
+        end
+
+        it { is_expected.to eq(expected_values) }
+
+        context 'when using descending order' do
+          let(:expected_values) { [5, 2] }
+
+          before do
+            params[:order] = :desc
+          end
+
+          it { is_expected.to eq(expected_values) }
+        end
+      end
+    end
+  end
 end