Improve discussion_line_code set and update notes_controller_spec

2 files changed, 7 insertions, 1 deletions

diff --git a/app/controllers/concerns/notes_actions.rb b/app/controllers/concerns/notes_actions.rb
index 688bf16cfea..d46a4dffa29 100644
--- a/app/controllers/concerns/notes_actions.rb
+++ b/app/controllers/concerns/notes_actions.rb
@@ -107,8 +107,9 @@ module NotesActions
 
             diff_discussion_html: diff_discussion_html(discussion),
             discussion_html: discussion_html(discussion),
-            discussion_line_code: discussion.line_code,
           )
+
+          attrs[:discussion_line_code] = discussion.line_code if discussion.is_a?(DiffDiscussion)
         end
       end
     else
diff --git a/spec/controllers/projects/notes_controller_spec.rb b/spec/controllers/projects/notes_controller_spec.rb
index 135fd6449ff..0081dcf599c 100644
--- a/spec/controllers/projects/notes_controller_spec.rb
+++ b/spec/controllers/projects/notes_controller_spec.rb
@@ -59,6 +59,7 @@ describe Projects::NotesController do
         expect(note_json[:id]).to eq(note.id)
         expect(note_json[:discussion_html]).not_to be_nil
         expect(note_json[:diff_discussion_html]).to be_nil
+        expect(note_json[:discussion_line_code]).to be_nil
       end
     end
 
@@ -74,6 +75,7 @@ describe Projects::NotesController do
         expect(note_json[:id]).to eq(note.id)
         expect(note_json[:discussion_html]).not_to be_nil
         expect(note_json[:diff_discussion_html]).not_to be_nil
+        expect(note_json[:discussion_line_code]).not_to be_nil
       end
     end
 
@@ -92,6 +94,7 @@ describe Projects::NotesController do
           expect(note_json[:id]).to eq(note.id)
           expect(note_json[:discussion_html]).not_to be_nil
           expect(note_json[:diff_discussion_html]).to be_nil
+          expect(note_json[:discussion_line_code]).to be_nil
         end
       end
 
@@ -104,6 +107,7 @@ describe Projects::NotesController do
           expect(note_json[:id]).to eq(note.id)
           expect(note_json[:discussion_html]).to be_nil
           expect(note_json[:diff_discussion_html]).to be_nil
+          expect(note_json[:discussion_line_code]).to be_nil
         end
       end
     end
@@ -120,6 +124,7 @@ describe Projects::NotesController do
         expect(note_json[:html]).not_to be_nil
         expect(note_json[:discussion_html]).to be_nil
         expect(note_json[:diff_discussion_html]).to be_nil
+        expect(note_json[:discussion_line_code]).to be_nil
       end
     end