1 files changed, 30 insertions, 4 deletions
diff --git a/spec/lib/gitlab/pagination/keyset/iterator_spec.rb b/spec/lib/gitlab/pagination/keyset/iterator_spec.rb
index d8e79287745..09cbca2c1cb 100644
--- a/spec/lib/gitlab/pagination/keyset/iterator_spec.rb
+++ b/spec/lib/gitlab/pagination/keyset/iterator_spec.rb
@@ -32,8 +32,11 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
     ])
   end
 
+  let(:iterator_params) { nil }
   let(:scope) { project.issues.reorder(custom_reorder) }
 
+  subject(:iterator) { described_class.new(**iterator_params) }
+
   shared_examples 'iterator examples' do
     describe '.each_batch' do
       it 'yields an ActiveRecord::Relation when a block is given' do
@@ -56,6 +59,29 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
         expect(count).to eq(9)
       end
 
+      it 'continues after the cursor' do
+        loaded_records = []
+        cursor = nil
+
+        # stopping the iterator after the first batch and storing the cursor
+        iterator.each_batch(of: 2) do |relation| # rubocop: disable Lint/UnreachableLoop
+          loaded_records.concat(relation.to_a)
+          record = loaded_records.last
+
+          cursor = custom_reorder.cursor_attributes_for_node(record)
+          break
+        end
+
+        expect(loaded_records).to eq(project.issues.order(custom_reorder).take(2))
+
+        new_iterator = described_class.new(**iterator_params.merge(cursor: cursor))
+        new_iterator.each_batch(of: 2) do |relation|
+          loaded_records.concat(relation.to_a)
+        end
+
+        expect(loaded_records).to eq(project.issues.order(custom_reorder))
+      end
+
       it 'allows updating of the yielded relations' do
         time = Time.current
 
@@ -73,7 +99,7 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
 
             iterator.each_batch(of: 2) { |rel| positions.concat(rel.pluck(:relative_position, :id)) }
 
-            expect(positions).to eq(project.issues.order_relative_position_asc.order(id: :asc).pluck(:relative_position, :id))
+            expect(positions).to eq(project.issues.reorder(::Gitlab::Database.nulls_last_order('relative_position', 'ASC')).order(id: :asc).pluck(:relative_position, :id))
           end
         end
 
@@ -85,7 +111,7 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
 
             iterator.each_batch(of: 2) { |rel| positions.concat(rel.pluck(:relative_position, :id)) }
 
-            expect(positions).to eq(project.issues.order_relative_position_desc.order(id: :desc).pluck(:relative_position, :id))
+            expect(positions).to eq(project.issues.reorder(::Gitlab::Database.nulls_first_order('relative_position', 'DESC')).order(id: :desc).pluck(:relative_position, :id))
           end
         end
 
@@ -131,13 +157,13 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
   end
 
   context 'when use_union_optimization is used' do
-    subject(:iterator) { described_class.new(scope: scope, use_union_optimization: true) }
+    let(:iterator_params) { { scope: scope, use_union_optimization: true } }
 
     include_examples 'iterator examples'
   end
 
   context 'when use_union_optimization is not used' do
-    subject(:iterator) { described_class.new(scope: scope, use_union_optimization: false) }
+    let(:iterator_params) { { scope: scope, use_union_optimization: false } }
 
     include_examples 'iterator examples'
   end