1 files changed, 40 insertions, 27 deletions

diff --git a/app/models/sent_notification.rb b/app/models/sent_notification.rb
index f4bcb49b34d..558e4c1a7eb 100644
--- a/app/models/sent_notification.rb
+++ b/app/models/sent_notification.rb
@@ -5,10 +5,11 @@ class SentNotification < ActiveRecord::Base
   belongs_to :noteable, polymorphic: true
   belongs_to :recipient, class_name: "User"
 
-  validates :project, :recipient, :reply_key, presence: true
-  validates :reply_key, uniqueness: true
+  validates :project, :recipient, presence: true
+  validates :reply_key, presence: true, uniqueness: true
   validates :noteable_id, presence: true, unless: :for_commit?
   validates :commit_id, presence: true, if: :for_commit?
+  validates :in_reply_to_discussion_id, format: { with: /\A\h{40}\z/, if: :discussion_note? }
   validate :note_valid
 
   after_save :keep_around_commit
@@ -34,23 +35,24 @@ class SentNotification < ActiveRecord::Base
       end
 
       attrs.reverse_merge!(
-        project:        noteable.project,
-        noteable_type:  noteable.class.name,
-        noteable_id:    noteable_id,
-        commit_id:      commit_id,
-        recipient_id:   recipient_id,
-        reply_key:      reply_key
+        project: noteable.project,
+        noteable_type: noteable.class.name,
+        noteable_id: noteable_id,
+        commit_id: commit_id,
+        recipient_id: recipient_id,
+        reply_key: reply_key
       )
 
       create(attrs)
     end
 
     def record_note(note, recipient_id, reply_key, attrs = {})
-      if note.diff_note?
-        attrs[:note_type] = note.type
+      attrs.merge!(
+        note_type: note.type,
+        in_reply_to_discussion_id: note.discussion_id
+      )
 
-        attrs.merge!(note.diff_attributes)
-      end
+      attrs.merge!(note.diff_attributes) if note.diff_note?
 
       record(note.noteable, recipient_id, reply_key, attrs)
     end
@@ -64,6 +66,10 @@ class SentNotification < ActiveRecord::Base
     noteable_type == "Commit"
   end
 
+  def discussion_note?
+    note_type == DiscussionNote.name
+  end
+
   def noteable
     if for_commit?
       project.commit(commit_id) rescue nil
@@ -89,31 +95,38 @@ class SentNotification < ActiveRecord::Base
     self.reply_key
   end
 
-  def note_attributes
-    {
-      project:        self.project,
-      author:         self.recipient,
-      type:           self.note_type,
-      noteable_type:  self.noteable_type,
-      noteable_id:    self.noteable_id,
-      commit_id:      self.commit_id,
-      line_code:      self.line_code,
-      position:       self.position.to_json
-    }
-  end
-
   def create_note(note)
     Notes::CreateService.new(
       self.project,
       self.recipient,
-      self.note_attributes.merge(note: note)
+      self.note_params.merge(note: note)
     ).execute
   end
 
   private
 
+  def note_params
+    {
+      project: self.project,
+      author: self.recipient,
+      type: self.note_type,
+      noteable_type: self.noteable_type,
+      noteable_id: self.noteable_id,
+      commit_id: self.commit_id,
+
+      # LegacyDiffNote
+      line_code: self.line_code,
+
+      # DiffNote
+      position: self.position.to_json,
+
+      # DiscussionNote
+      in_reply_to_discussion_id: self.in_reply_to_discussion_id
+    }
+  end
+
   def note_valid
-    Note.new(note_attributes.merge(note: "Test")).valid?
+    Note.new(note_params.merge(note: "Test")).valid?
   end
 
   def keep_around_commit