1 files changed, 50 insertions, 5 deletions
diff --git a/spec/models/merge_request_diff_spec.rb b/spec/models/merge_request_diff_spec.rb
index 4075eb96fc2..3741e01e99a 100644
--- a/spec/models/merge_request_diff_spec.rb
+++ b/spec/models/merge_request_diff_spec.rb
@@ -418,8 +418,8 @@ RSpec.describe MergeRequestDiff do
       shared_examples_for 'fetching full diffs' do
         it 'returns diffs from repository comparison' do
           expect_next_instance_of(Compare) do |comparison|
-            expect(comparison).to receive(:diffs_in_batch)
-              .with(1, 10, diff_options: diff_options)
+            expect(comparison).to receive(:diffs)
+              .with(diff_options)
               .and_call_original
           end
 
@@ -448,13 +448,13 @@ RSpec.describe MergeRequestDiff do
         end
 
         it_behaves_like 'fetching full diffs'
-      end
 
-      context 'when diff_options include ignore_whitespace_change' do
-        it_behaves_like 'fetching full diffs' do
+        context 'when diff_options include ignore_whitespace_change' do
           let(:diff_options) do
             { ignore_whitespace_change: true }
           end
+
+          it_behaves_like 'fetching full diffs'
         end
       end
 
@@ -485,6 +485,51 @@ RSpec.describe MergeRequestDiff do
             'files/whitespace'
           ])
         end
+
+        context 'when diff_options include ignore_whitespace_change' do
+          let(:diff_options) do
+            { ignore_whitespace_change: true }
+          end
+
+          it 'returns a Gitlab::Diff::FileCollection::Compare with paginated diffs' do
+            diffs = diff_with_commits.diffs_in_batch(1, 10, diff_options: diff_options)
+
+            expect(diffs).to be_a(Gitlab::Diff::FileCollection::Compare)
+            expect(diffs.diff_files.size).to eq 10
+            expect(diffs.pagination_data).to eq(current_page: 1, next_page: 2, total_pages: 2)
+          end
+
+          it 'returns an empty MergeRequestBatch with empty pagination data when the batch is empty' do
+            diffs = diff_with_commits.diffs_in_batch(3, 10, diff_options: diff_options)
+
+            expect(diffs).to be_a(Gitlab::Diff::FileCollection::MergeRequestDiffBatch)
+            expect(diffs.diff_files.size).to eq 0
+            expect(diffs.pagination_data).to eq(current_page: nil, next_page: nil, total_pages: nil)
+          end
+
+          context 'with gradual load enabled' do
+            before do
+              stub_feature_flags(diffs_gradual_load: true)
+            end
+
+            it 'returns pagination data from MergeRequestDiffBatch' do
+              diffs = diff_with_commits.diffs_in_batch(1, 10, diff_options: diff_options)
+              file_count = diff_with_commits.merge_request_diff_files.count
+
+              expect(diffs).to be_a(Gitlab::Diff::FileCollection::Compare)
+              expect(diffs.diff_files.size).to eq 10
+              expect(diffs.pagination_data).to eq(current_page: nil, next_page: nil, total_pages: file_count)
+            end
+
+            it 'returns an empty MergeRequestBatch with empty pagination data when the batch is empty' do
+              diffs = diff_with_commits.diffs_in_batch(30, 10, diff_options: diff_options)
+
+              expect(diffs).to be_a(Gitlab::Diff::FileCollection::MergeRequestDiffBatch)
+              expect(diffs.diff_files.size).to eq 0
+              expect(diffs.pagination_data).to eq(current_page: nil, next_page: nil, total_pages: nil)
+            end
+          end
+        end
       end
     end