1 files changed, 27 insertions, 6 deletions
diff --git a/spec/models/merge_request_diff_spec.rb b/spec/models/merge_request_diff_spec.rb
index 90cce826b6c..47e8f04e728 100644
--- a/spec/models/merge_request_diff_spec.rb
+++ b/spec/models/merge_request_diff_spec.rb
@@ -52,9 +52,9 @@ describe MergeRequestDiff do
 
     context 'when it was not cleaned by the system' do
       it 'returns persisted diffs' do
-        expect(diff).to receive(:load_diffs)
+        expect(diff).to receive(:load_diffs).and_call_original
 
-        diff.diffs
+        diff.diffs.diff_files
       end
     end
 
@@ -76,19 +76,19 @@ describe MergeRequestDiff do
       end
 
       it 'returns persisted diffs if cannot compare with diff refs' do
-        expect(diff).to receive(:load_diffs)
+        expect(diff).to receive(:load_diffs).and_call_original
 
         diff.update!(head_commit_sha: 'invalid-sha')
 
-        diff.diffs
+        diff.diffs.diff_files
       end
 
       it 'returns persisted diffs if diff refs does not exist' do
-        expect(diff).to receive(:load_diffs)
+        expect(diff).to receive(:load_diffs).and_call_original
 
         diff.update!(start_commit_sha: nil, base_commit_sha: nil)
 
-        diff.diffs
+        diff.diffs.diff_files
       end
     end
   end
@@ -211,4 +211,25 @@ describe MergeRequestDiff do
       expect(diff_with_commits.commits_count).to eq(29)
     end
   end
+
+  describe '#commits_by_shas' do
+    let(:commit_shas) { diff_with_commits.commit_shas }
+
+    it 'returns empty if no SHAs were provided' do
+      expect(diff_with_commits.commits_by_shas([])).to be_empty
+    end
+
+    it 'returns one SHA' do
+      commits = diff_with_commits.commits_by_shas([commit_shas.first, Gitlab::Git::BLANK_SHA])
+
+      expect(commits.count).to eq(1)
+    end
+
+    it 'returns all matching SHAs' do
+      commits = diff_with_commits.commits_by_shas(commit_shas)
+
+      expect(commits.count).to eq(commit_shas.count)
+      expect(commits.map(&:sha)).to match_array(commit_shas)
+    end
+  end
 end