Fetch projects from subgroups in ProjectsGrade

ee/app/models/vulnerabilities/projects_grade.rb

@@ -2,21 +2,23 @@
 
 module Vulnerabilities
   class ProjectsGrade
-    attr_reader :vulnerable, :grade, :project_ids
+    attr_reader :vulnerable, :grade, :project_ids, :include_subgroups
 
     # project_ids can contain IDs from projects that do not belong to vulnerable, they will be filtered out in `projects` method
-    def initialize(vulnerable, letter_grade, project_ids = [])
+    def initialize(vulnerable, letter_grade, project_ids = [], include_subgroups: false)
       @vulnerable = vulnerable
       @grade = letter_grade
       @project_ids = project_ids
+      @include_subgroups = include_subgroups
     end
 
     delegate :count, to: :projects
 
     def projects
-      return vulnerable.projects.none if project_ids.blank?
+      return Project.none if project_ids.blank?
 
-      vulnerable.projects.with_vulnerability_statistics.inc_routes.where(id: project_ids)
+      projects = include_subgroups ? vulnerable.all_projects : vulnerable.projects
+      projects.with_vulnerability_statistics.inc_routes.where(id: project_ids)
     end
 
     def self.grades_for(vulnerables, include_subgroups: false)
@@ -28,7 +30,7 @@ module Vulnerabilities
         .select(:letter_grade, 'array_agg(project_id) project_ids')
         .then do |statistics|
           vulnerables.each_with_object({}) do |vulnerable, hash|
-            hash[vulnerable] = statistics.map { |statistic| new(vulnerable, statistic.letter_grade, statistic.project_ids) }
+            hash[vulnerable] = statistics.map { |statistic| new(vulnerable, statistic.letter_grade, statistic.project_ids, include_subgroups: include_subgroups) }
           end
         end
     end

ee/spec/models/vulnerabilities/projects_grade_spec.rb

@@ -98,17 +98,34 @@ RSpec.describe Vulnerabilities::ProjectsGrade do
   end
 
   describe '#projects' do
-    let(:projects_grade) { described_class.new(group, 1, [project_3.id, project_4.id]) }
-    let(:expected_projects) { [project_3, project_4] }
+    let(:projects_grade) { described_class.new(group, 1, [project_3.id, project_4.id, project_6.id], include_subgroups: include_subgroups) }
 
     subject(:projects) { projects_grade.projects }
 
-    it { is_expected.to match_array(expected_projects) }
+    context 'when include_subgroups is set to false' do
+      let(:include_subgroups) { false }
+      let(:expected_projects) { [project_3, project_4] }
+
+      it { is_expected.to match_array(expected_projects) }
+
+      it 'preloads vulnerability statistic once for whole collection' do
+        control_count = ActiveRecord::QueryRecorder.new { described_class.new(group, 1, [project_3.id]).projects.map(&:vulnerability_statistic) }.count
 
-    it 'preloads vulnerability statistic once for whole collection' do
-      control_count = ActiveRecord::QueryRecorder.new { described_class.new(group, 1, [project_3.id]).projects.map(&:vulnerability_statistic) }.count
+        expect { described_class.new(group, 1, [project_3.id, project_4.id]).projects.map(&:vulnerability_statistic) }.not_to exceed_query_limit(control_count)
+      end
+    end
 
-      expect { described_class.new(group, 1, [project_3.id, project_4.id]).projects.map(&:vulnerability_statistic) }.not_to exceed_query_limit(control_count)
+    context 'when include_subgroups is set to true' do
+      let(:include_subgroups) { true }
+      let(:expected_projects) { [project_3, project_4, project_6] }
+
+      it { is_expected.to match_array(expected_projects) }
+
+      it 'preloads vulnerability statistic once for whole collection' do
+        control_count = ActiveRecord::QueryRecorder.new { described_class.new(group, 1, [project_3.id]).projects.map(&:vulnerability_statistic) }.count
+
+        expect { described_class.new(group, 1, [project_3.id, project_4.id]).projects.map(&:vulnerability_statistic) }.not_to exceed_query_limit(control_count)
+      end
     end
   end