Ensure milestones passed to GlobalMilestone is an ActiveRecord::Relation

2 files changed, 7 insertions, 15 deletions

diff --git a/app/models/global_milestone.rb b/app/models/global_milestone.rb
index 8f38a2b6254..3e10bed19c2 100644
--- a/app/models/global_milestone.rb
+++ b/app/models/global_milestone.rb
@@ -8,7 +8,8 @@ class GlobalMilestone
     milestones = milestones.group_by(&:title)
 
     milestones.map do |title, milestones|
-      new(title, milestones)
+      milestones_relation = Milestone.where(id: milestones.map(&:id))
+      new(title, milestones_relation)
     end
   end
 
@@ -61,11 +62,11 @@ class GlobalMilestone
   end
 
   def participants
-    @participants ||= milestones_relation.includes(:participants).map(&:participants).flatten.compact.uniq
+    @participants ||= milestones.includes(:participants).map(&:participants).flatten.compact.uniq
   end
 
   def labels
-    @labels ||= GlobalLabel.build_collection(milestones_relation.includes(:labels).map(&:labels).flatten)
+    @labels ||= GlobalLabel.build_collection(milestones.includes(:labels).map(&:labels).flatten)
                            .sort_by!(&:title)
   end
 
@@ -89,14 +90,4 @@ class GlobalMilestone
       end
     end
   end
-
-  private
-
-  def milestones_relation
-    @milestones_relation ||= if milestones.is_a?(ActiveRecord::Relation)
-                               milestones
-                             else
-                               Milestone.where(id: milestones.map(&:id))
-                             end
-  end
 end
diff --git a/spec/models/global_milestone_spec.rb b/spec/models/global_milestone_spec.rb
index 92e0f7f27ce..dd033480527 100644
--- a/spec/models/global_milestone_spec.rb
+++ b/spec/models/global_milestone_spec.rb
@@ -50,8 +50,9 @@ describe GlobalMilestone, models: true do
           milestone1_project2,
           milestone1_project3,
         ]
+      milestones_relation = Milestone.where(id: milestones.map(&:id))
 
-      @global_milestone = GlobalMilestone.new(milestone1_project1.title, milestones)
+      @global_milestone = GlobalMilestone.new(milestone1_project1.title, milestones_relation)
     end
 
     it 'has exactly one group milestone' do
@@ -67,7 +68,7 @@ describe GlobalMilestone, models: true do
     let(:milestone) { create(:milestone, title: "git / test", project: project1) }
 
     it 'strips out slashes and spaces' do
-      global_milestone = GlobalMilestone.new(milestone.title, [milestone])
+      global_milestone = GlobalMilestone.new(milestone.title, Milestone.where(id: milestone.id))
 
       expect(global_milestone.safe_title).to eq('git-test')
     end