1 files changed, 24 insertions, 10 deletions
diff --git a/app/finders/notes_finder.rb b/app/finders/notes_finder.rb
index a653a6d59c6..fc86acd9329 100644
--- a/app/finders/notes_finder.rb
+++ b/app/finders/notes_finder.rb
@@ -1,27 +1,41 @@
 class NotesFinder
   FETCH_OVERLAP = 5.seconds
 
-  def execute(project, current_user, params)
+  attr_accessor :project, :current_user, :params
+
+  def initialize(project, current_user, params)
+    @project, @current_user, @params = project, current_user, params
+  end
+
+  def execute
+    notes =
+      if target.respond_to?(:related_notes)
+        target.related_notes
+      else
+        target.notes
+      end
+
+    last_fetched_at = Time.at(params.fetch(:last_fetched_at, 0).to_i)
+    # Use overlapping intervals to avoid worrying about race conditions
+    notes.inc_author.where('updated_at > ?', last_fetched_at - FETCH_OVERLAP).fresh
+  end
+
+  def target
     target_type = params[:target_type]
     target_id   = params[:target_id]
-    # Default to 0 to remain compatible with old clients
-    last_fetched_at = Time.at(params.fetch(:last_fetched_at, 0).to_i)
 
     notes =
       case target_type
       when "commit"
-        project.notes.for_commit_id(target_id).non_diff_notes
+        project.commit(target_id)
       when "issue"
-        IssuesFinder.new(current_user, project_id: project.id).find(target_id).notes.inc_author
+        IssuesFinder.new(current_user, project_id: project.id).find(target_id)
       when "merge_request"
-        project.merge_requests.find(target_id).mr_and_commit_notes.inc_author
+        project.merge_requests.find(target_id)
       when "snippet", "project_snippet"
-        project.snippets.find(target_id).notes
+        project.snippets.find(target_id)
       else
         raise 'invalid target_type'
       end
-
-    # Use overlapping intervals to avoid worrying about race conditions
-    notes.where('updated_at > ?', last_fetched_at - FETCH_OVERLAP).fresh
   end
 end