Refactor discussions/notes code

2 files changed, 6 insertions, 2 deletions

diff --git a/app/controllers/concerns/issuable_actions.rb b/app/controllers/concerns/issuable_actions.rb
index a21e658fda1..0379f76fc3d 100644
--- a/app/controllers/concerns/issuable_actions.rb
+++ b/app/controllers/concerns/issuable_actions.rb
@@ -88,11 +88,15 @@ module IssuableActions
 
     discussions = Discussion.build_collection(notes, issuable)
 
-    render json: DiscussionSerializer.new(project: project, noteable: issuable, current_user: current_user).represent(discussions, context: self)
+    render json: discussion_serializer.represent(discussions, context: self)
   end
 
   private
 
+  def discussion_serializer
+    DiscussionSerializer.new(project: project, noteable: issuable, current_user: current_user, note_entity: ProjectNoteEntity)
+  end
+
   def recaptcha_check_if_spammable(should_redirect = true, &block)
     return yield unless issuable.is_a? Spammable
 
diff --git a/app/controllers/concerns/notes_actions.rb b/app/controllers/concerns/notes_actions.rb
index 03ed5b5310b..839cac3687c 100644
--- a/app/controllers/concerns/notes_actions.rb
+++ b/app/controllers/concerns/notes_actions.rb
@@ -212,7 +212,7 @@ module NotesActions
   end
 
   def note_serializer
-    NoteSerializer.new(project: project, noteable: noteable, current_user: current_user)
+    ProjectNoteSerializer.new(project: project, noteable: noteable, current_user: current_user)
   end
 
   def note_project