1 files changed, 64 insertions, 27 deletions
diff --git a/spec/lib/gitlab/git/diff_collection_spec.rb b/spec/lib/gitlab/git/diff_collection_spec.rb
index 122c93dcd69..3494f0cc98d 100644
--- a/spec/lib/gitlab/git/diff_collection_spec.rb
+++ b/spec/lib/gitlab/git/diff_collection_spec.rb
@@ -6,18 +6,18 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
       iterator,
       max_files: max_files,
       max_lines: max_lines,
-      all_diffs: all_diffs,
-      no_collapse: no_collapse
+      limits: limits,
+      expanded: expanded
     )
   end
-  let(:iterator) { Array.new(file_count, fake_diff(line_length, line_count)) }
+  let(:iterator) { MutatingConstantIterator.new(file_count, fake_diff(line_length, line_count)) }
   let(:file_count) { 0 }
   let(:line_length) { 1 }
   let(:line_count) { 1 }
   let(:max_files) { 10 }
   let(:max_lines) { 100 }
-  let(:all_diffs) { false }
-  let(:no_collapse) { true }
+  let(:limits) { true }
+  let(:expanded) { true }
 
   describe '#to_a' do
     subject { super().to_a }
@@ -64,10 +64,18 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
           subject { super().real_size }
           it { is_expected.to eq('3') }
         end
-        it { expect(subject.size).to eq(3) }
+
+        describe '#size' do
+          it { expect(subject.size).to eq(3) }
+
+          it 'does not change after peeking' do
+            subject.any?
+            expect(subject.size).to eq(3)
+          end
+        end
 
         context 'when limiting is disabled' do
-          let(:all_diffs) { true }
+          let(:limits) { false }
 
           describe '#overflow?' do
             subject { super().overflow? }
@@ -83,7 +91,15 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
             subject { super().real_size }
             it { is_expected.to eq('3') }
           end
-          it { expect(subject.size).to eq(3) }
+
+          describe '#size' do
+            it { expect(subject.size).to eq(3) }
+
+            it 'does not change after peeking' do
+              subject.any?
+              expect(subject.size).to eq(3)
+            end
+          end
         end
       end
 
@@ -107,7 +123,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
         it { expect(subject.size).to eq(0) }
 
         context 'when limiting is disabled' do
-          let(:all_diffs) { true }
+          let(:limits) { false }
 
           describe '#overflow?' do
             subject { super().overflow? }
@@ -151,7 +167,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
         it { expect(subject.size).to eq(10) }
 
         context 'when limiting is disabled' do
-          let(:all_diffs) { true }
+          let(:limits) { false }
 
           describe '#overflow?' do
             subject { super().overflow? }
@@ -191,7 +207,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
         it { expect(subject.size).to eq(3) }
 
         context 'when limiting is disabled' do
-          let(:all_diffs) { true }
+          let(:limits) { false }
 
           describe '#overflow?' do
             subject { super().overflow? }
@@ -257,7 +273,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
       it { expect(subject.size).to eq(9) }
 
       context 'when limiting is disabled' do
-        let(:all_diffs) { true }
+        let(:limits) { false }
 
         describe '#overflow?' do
           subject { super().overflow? }
@@ -309,8 +325,8 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
       end
 
       it 'yields Diff instances even when they are too large' do
-        expect { |b| collection.each(&b) }.
-          to yield_with_args(an_instance_of(Gitlab::Git::Diff))
+        expect { |b| collection.each(&b) }
+          .to yield_with_args(an_instance_of(Gitlab::Git::Diff))
       end
 
       it 'prunes diffs that are too large' do
@@ -325,14 +341,15 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
     end
 
     context 'when diff is quite large will collapse by default' do
-      let(:iterator) { [{ diff: 'a' * 20480 }] }
+      let(:iterator) { [{ diff: 'a' * (Gitlab::Git::Diff.collapse_limit + 1) }] }
+      let(:max_files) { 100 }
 
       context 'when no collapse is set' do
-        let(:no_collapse) { true }
+        let(:expanded) { true }
 
         it 'yields Diff instances even when they are quite big' do
-          expect { |b| subject.each(&b) }.
-            to yield_with_args(an_instance_of(Gitlab::Git::Diff))
+          expect { |b| subject.each(&b) }
+            .to yield_with_args(an_instance_of(Gitlab::Git::Diff))
         end
 
         it 'does not prune diffs' do
@@ -347,11 +364,11 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
       end
 
       context 'when no collapse is unset' do
-        let(:no_collapse) { false }
+        let(:expanded) { false }
 
         it 'yields Diff instances even when they are quite big' do
-          expect { |b| subject.each(&b) }.
-            to yield_with_args(an_instance_of(Gitlab::Git::Diff))
+          expect { |b| subject.each(&b) }
+            .to yield_with_args(an_instance_of(Gitlab::Git::Diff))
         end
 
         it 'prunes diffs that are quite big' do
@@ -367,7 +384,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
         context 'when go over safe limits on files' do
           let(:iterator) { [fake_diff(1, 1)] * 4 }
 
-          before(:each) do
+          before do
             stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS', { max_files: 2, max_lines: max_lines })
           end
 
@@ -392,7 +409,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
             ]
           end
 
-          before(:each) do
+          before do
             stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS', { max_files: max_files, max_lines: 80 })
           end
 
@@ -417,7 +434,7 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
             ]
           end
 
-          before(:each) do
+          before do
             stub_const('Gitlab::Git::DiffCollection::DEFAULT_LIMITS', { max_files: max_files, max_lines: 80 })
           end
 
@@ -434,11 +451,11 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
       end
 
       context 'when limiting is disabled' do
-        let(:all_diffs) { true }
+        let(:limits) { false }
 
         it 'yields Diff instances even when they are quite big' do
-          expect { |b| subject.each(&b) }.
-            to yield_with_args(an_instance_of(Gitlab::Git::Diff))
+          expect { |b| subject.each(&b) }
+            .to yield_with_args(an_instance_of(Gitlab::Git::Diff))
         end
 
         it 'does not prune diffs' do
@@ -457,4 +474,24 @@ describe Gitlab::Git::DiffCollection, seed_helper: true do
   def fake_diff(line_length, line_count)
     { 'diff' => "#{'a' * line_length}\n" * line_count }
   end
+
+  class MutatingConstantIterator
+    include Enumerable
+
+    def initialize(count, value)
+      @count = count
+      @value = value
+    end
+
+    def each
+      return enum_for(:each) unless block_given?
+
+      loop do
+        break if @count.zero?
+        # It is critical to decrement before yielding. We may never reach the lines after 'yield'.
+        @count -= 1
+        yield @value
+      end
+    end
+  end
 end