1 files changed, 36 insertions, 10 deletions

diff --git a/spec/lib/gitlab/git/diff_collection_spec.rb b/spec/lib/gitlab/git/diff_collection_spec.rb
index b202015464f..1a3c332a21b 100644
--- a/spec/lib/gitlab/git/diff_collection_spec.rb
+++ b/spec/lib/gitlab/git/diff_collection_spec.rb
@@ -9,8 +9,11 @@ RSpec.describe Gitlab::Git::DiffCollection, :seed_helper do
     MutatingConstantIterator.class_eval do
       include Enumerable
 
+      attr_reader :size
+
       def initialize(count, value)
         @count = count
+        @size  = count
         @value = value
       end
 
@@ -517,21 +520,39 @@ RSpec.describe Gitlab::Git::DiffCollection, :seed_helper do
             .to yield_with_args(an_instance_of(Gitlab::Git::Diff))
         end
 
-        it 'prunes diffs that are quite big' do
-          diff = nil
+        context 'single-file collections' do
+          it 'does not prune diffs' do
+            diff = nil
 
-          subject.each do |d|
-            diff = d
+            subject.each do |d|
+              diff = d
+            end
+
+            expect(diff.diff).not_to eq('')
           end
+        end
+
+        context 'multi-file collections' do
+          let(:iterator) { [{ diff: 'b' }, { diff: 'a' * 20480 }]}
+
+          it 'prunes diffs that are quite big' do
+            diff = nil
 
-          expect(diff.diff).to eq('')
+            subject.each do |d|
+              diff = d
+            end
+
+            expect(diff.diff).to eq('')
+          end
         end
 
         context 'when go over safe limits on files' do
           let(:iterator) { [fake_diff(1, 1)] * 4 }
 
           before do
-            stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS', { max_files: 2, max_lines: max_lines })
+            allow(Gitlab::Git::DiffCollection)
+              .to receive(:default_limits)
+              .and_return({ max_files: 2, max_lines: max_lines })
           end
 
           it 'prunes diffs by default even little ones' do
@@ -556,7 +577,9 @@ RSpec.describe Gitlab::Git::DiffCollection, :seed_helper do
           end
 
           before do
-            stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS', { max_files: max_files, max_lines: 80 })
+            allow(Gitlab::Git::DiffCollection)
+              .to receive(:default_limits)
+              .and_return({ max_files: max_files, max_lines: 80 })
           end
 
           it 'prunes diffs by default even little ones' do
@@ -581,7 +604,9 @@ RSpec.describe Gitlab::Git::DiffCollection, :seed_helper do
           end
 
           before do
-            stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS', { max_files: max_files, max_lines: 80 })
+            allow(Gitlab::Git::DiffCollection)
+              .to receive(:default_limits)
+              .and_return({ max_files: max_files, max_lines: 80 })
           end
 
           it 'prunes diffs by default even little ones' do
@@ -665,8 +690,9 @@ RSpec.describe Gitlab::Git::DiffCollection, :seed_helper do
         end
 
         before do
-          stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS',
-                     { max_files: max_files, max_lines: 80 })
+          allow(Gitlab::Git::DiffCollection)
+            .to receive(:default_limits)
+            .and_return({ max_files: max_files, max_lines: 80 })
         end
 
         it 'considers size of diffs before the offset for prunning' do