2 files changed, 87 insertions, 3 deletions

diff --git a/spec/models/diff_note_spec.rb b/spec/models/diff_note_spec.rb
index 4aa9ec789a3..eb0a3e9e0d3 100644
--- a/spec/models/diff_note_spec.rb
+++ b/spec/models/diff_note_spec.rb
@@ -3,7 +3,7 @@ require 'spec_helper'
 describe DiffNote do
   include RepoHelpers
 
-  let(:merge_request) { create(:merge_request) }
+  let!(:merge_request) { create(:merge_request) }
   let(:project) { merge_request.project }
   let(:commit) { project.commit(sample_commit.id) }
 
@@ -98,14 +98,14 @@ describe DiffNote do
       diff_line = subject.diff_line
 
       expect(diff_line.added?).to be true
-      expect(diff_line.new_line).to eq(position.new_line)
+      expect(diff_line.new_line).to eq(position.formatter.new_line)
       expect(diff_line.text).to eq("+    vars = {")
     end
   end
 
   describe "#line_code" do
     it "returns the correct line code" do
-      line_code = Gitlab::Diff::LineCode.generate(position.file_path, position.new_line, 15)
+      line_code = Gitlab::Diff::LineCode.generate(position.file_path, position.formatter.new_line, 15)
 
       expect(subject.line_code).to eq(line_code)
     end
@@ -255,4 +255,38 @@ describe DiffNote do
       end
     end
   end
+
+  describe "image diff notes" do
+    let(:path) { "files/images/any_image.png" }
+
+    let!(:position) do
+      Gitlab::Diff::Position.new(
+        old_path: path,
+        new_path: path,
+        width: 10,
+        height: 10,
+        x: 1,
+        y: 1,
+        diff_refs: merge_request.diff_refs,
+        position_type: "image"
+      )
+    end
+
+    describe "validations" do
+      subject { build(:diff_note_on_merge_request, project: project, position: position, noteable: merge_request) }
+
+      it { is_expected.not_to validate_presence_of(:line_code) }
+
+      it "does not validate diff line" do
+        diff_line = subject.diff_line
+
+        expect(diff_line).to be nil
+        expect(subject).to be_valid
+      end
+    end
+
+    it "returns true for on_image?" do
+      expect(subject.on_image?).to be_truthy
+    end
+  end
 end
diff --git a/spec/models/note_spec.rb b/spec/models/note_spec.rb
index b214074fdce..1ecb50586c7 100644
--- a/spec/models/note_spec.rb
+++ b/spec/models/note_spec.rb
@@ -314,6 +314,56 @@ describe Note do
         expect(subject[active_diff_note1.line_code].first.id).to eq(active_diff_note1.discussion_id)
         expect(subject[active_diff_note3.line_code].first.id).to eq(active_diff_note3.discussion_id)
       end
+
+      context 'with image discussions' do
+        let(:merge_request2) { create(:merge_request_with_diffs, :with_image_diffs, source_project: project, title: "Added images and changes") }
+        let(:image_path) { "files/images/ee_repo_logo.png" }
+        let(:text_path) { "bar/branch-test.txt" }
+        let!(:image_note) { create(:diff_note_on_merge_request, project: project, noteable: merge_request2, position: image_position) }
+        let!(:text_note) { create(:diff_note_on_merge_request, project: project, noteable: merge_request2, position: text_position) }
+
+        let(:image_position) do
+          Gitlab::Diff::Position.new(
+            old_path: image_path,
+            new_path: image_path,
+            width: 100,
+            height: 100,
+            x: 1,
+            y: 1,
+            position_type: "image",
+            diff_refs: merge_request2.diff_refs
+          )
+        end
+
+        let(:text_position) do
+          Gitlab::Diff::Position.new(
+            old_path: text_path,
+            new_path: text_path,
+            old_line: nil,
+            new_line: 2,
+            position_type: "text",
+            diff_refs: merge_request2.diff_refs
+          )
+        end
+
+        it "groups image discussions by file identifier" do
+          diff_discussion = DiffDiscussion.new([image_note])
+
+          discussions = merge_request2.notes.grouped_diff_discussions
+
+          expect(discussions.size).to eq(2)
+          expect(discussions[image_note.diff_file.new_path]).to include(diff_discussion)
+        end
+
+        it "groups text discussions by line code" do
+          diff_discussion = DiffDiscussion.new([text_note])
+
+          discussions = merge_request2.notes.grouped_diff_discussions
+
+          expect(discussions.size).to eq(2)
+          expect(discussions[text_note.line_code]).to include(diff_discussion)
+        end
+      end
     end
 
     context 'diff discussions for older diff refs' do