Fix userNotesCount aggregate

This change fixes problem where 2 aggregates were using same key to
store its state. This change also renames lazy_aggregate to more
specific name to prevent similar situation to occur in the future.

ee/app/graphql/ee/gitlab_schema.rb

@@ -7,8 +7,8 @@ module EE
     prepended do
       lazy_resolve ::Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate, :epic_aggregate
       lazy_resolve ::Gitlab::Graphql::Aggregations::Issues::LazyBlockAggregate, :block_aggregate
-      lazy_resolve ::Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggregate, :aggregate
-      lazy_resolve ::Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCountAggregate, :user_notes_count_aggregate
+      lazy_resolve ::Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggregate, :execute
+      lazy_resolve ::Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCountAggregate, :execute
     end
   end
 end

ee/lib/gitlab/graphql/aggregations/vulnerabilities/lazy_user_notes_count_aggregate.rb

@@ -12,7 +12,7 @@ module Gitlab
 
             # Initialize the loading state for this query,
             # or get the previously-initiated state
-            @lazy_state = query_ctx[:lazy_aggregate] ||= {
+            @lazy_state = query_ctx[:lazy_user_notes_count_aggregate] ||= {
               pending_vulnerability_ids: Set.new,
               loaded_objects: {}
             }
@@ -21,7 +21,7 @@ module Gitlab
           end
 
           # Return the loaded record, hitting the database if needed
-          def user_notes_count_aggregate
+          def execute
             # Check if the record was already loaded
             if @lazy_state[:pending_vulnerability_ids].present?
               load_records_into_loaded_objects

ee/lib/gitlab/graphql/aggregations/vulnerability_statistics/lazy_aggregate.rb

@@ -12,7 +12,7 @@ module Gitlab
 
             # Initialize the loading state for this query,
             # or get the previously-initiated state
-            @lazy_state = query_ctx[:lazy_aggregate] ||= {
+            @lazy_state = query_ctx[:lazy_vulnerability_statistics_aggregate] ||= {
               pending_vulnerables: Set.new,
               loaded_objects: {}
             }
@@ -21,7 +21,7 @@ module Gitlab
           end
 
           # Return the loaded record, hitting the database if needed
-          def aggregate
+          def execute
             # Check if the record was already loaded
             if @lazy_state[:pending_vulnerables].present?
               load_records_into_loaded_objects

ee/spec/factories/vulnerabilities.rb

@@ -57,6 +57,20 @@ FactoryBot.define do
       end
     end
 
+    trait :with_notes do
+      transient do
+        notes_count { 3 }
+      end
+
+      after(:create) do |vulnerability, evaluator|
+        create_list(
+          :note_on_vulnerability,
+          evaluator.notes_count,
+          noteable: vulnerability,
+          project: vulnerability.project)
+      end
+    end
+
     trait :with_findings do
       after(:build) do |vulnerability|
         findings_with_solution = build_list(

ee/spec/lib/gitlab/graphql/aggregations/vulnerabilities/lazy_user_notes_count_aggregate_spec.rb

@@ -16,9 +16,16 @@ RSpec.describe Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCoun
       expect(subject.lazy_state[:pending_vulnerability_ids]).to match [vulnerability.id]
       expect(subject.vulnerability).to match vulnerability
     end
+
+    it 'uses lazy_user_notes_count_aggregate to collect aggregates' do
+      subject = described_class.new({ lazy_user_notes_count_aggregate: { pending_vulnerability_ids: [10, 20, 30].to_set, loaded_objects: {} } }, vulnerability)
+
+      expect(subject.lazy_state[:pending_vulnerability_ids]).to match [10, 20, 30, vulnerability.id]
+      expect(subject.vulnerability).to match vulnerability
+    end
   end
 
-  describe '#user_notes_count_aggregate' do
+  describe '#execute' do
     subject { described_class.new(query_ctx, vulnerability) }
 
     before do
@@ -33,7 +40,7 @@ RSpec.describe Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCoun
       it 'does not make the query again' do
         expect(::Note).not_to receive(:user)
 
-        subject.user_notes_count_aggregate
+        subject.execute
       end
     end
 
@@ -57,11 +64,11 @@ RSpec.describe Gitlab::Graphql::Aggregations::Vulnerabilities::LazyUserNotesCoun
       it 'makes the query' do
         expect(::Note).to receive_message_chain(:user, :count_for_vulnerability_id).with([vulnerability.id, other_vulnerability.id])
 
-        subject.user_notes_count_aggregate
+        subject.execute
       end
 
       it 'clears the pending IDs' do
-        subject.user_notes_count_aggregate
+        subject.execute
 
         expect(subject.lazy_state[:pending_vulnerability_ids]).to be_empty
       end

ee/spec/lib/gitlab/graphql/aggregations/vulnerability_statistics/lazy_aggregate_spec.rb

@@ -20,7 +20,7 @@ RSpec.describe Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggre
     end
   end
 
-  describe '#aggregate' do
+  describe '#execute' do
     subject { described_class.new(query_ctx, vulnerable) }
 
     before do
@@ -35,7 +35,7 @@ RSpec.describe Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggre
       it 'does not make the query again' do
         expect(::Vulnerabilities::ProjectsGrade).not_to receive(:grades_for)
 
-        subject.aggregate
+        subject.execute
       end
     end
 
@@ -59,11 +59,11 @@ RSpec.describe Gitlab::Graphql::Aggregations::VulnerabilityStatistics::LazyAggre
       it 'makes the query' do
         expect(::Vulnerabilities::ProjectsGrade).to receive(:grades_for).with([vulnerable, other_vulnerable])
 
-        subject.aggregate
+        subject.execute
       end
 
       it 'clears the pending IDs' do
-        subject.aggregate
+        subject.execute
 
         expect(subject.lazy_state[:pending_vulnerables]).to be_empty
       end

ee/spec/requests/api/graphql/group_query_spec.rb

@@ -43,5 +43,158 @@ RSpec.describe 'getting group information' do
         expect(graphql_data['group']['id']).to eq(group.to_global_id.to_s)
       end
     end
+
+    context 'when loading vulnerabilityGrades alongside with Vulnerability.userNotesCount' do
+      let_it_be(:private_group) { create(:group, :private) }
+      let_it_be(:public_group) { create(:group, :public) }
+
+      let(:fields_public_group) do
+        <<~QUERY
+        vulnerabilityGrades {
+          grade
+          count
+          projects {
+            nodes {
+              vulnerabilities {
+                nodes {
+                  id
+                  userNotesCount
+                }
+              }
+            }
+          }
+        }
+        QUERY
+      end
+
+      let(:fields_private_group) do
+        <<~QUERY
+        vulnerabilityGrades {
+          grade
+          count
+          projects {
+            nodes {
+              confirmedVulnerabilities: vulnerabilities(state: CONFIRMED) {
+                nodes {
+                  id
+                  userNotesCount
+                }
+              }
+              dismissedVulnerabilities: vulnerabilities(state: DISMISSED) {
+                nodes {
+                  id
+                  userNotesCount
+                }
+              }
+            }
+          }
+        }
+        QUERY
+      end
+
+      let(:queries) do
+        [
+          { query: graphql_query_for('group', { 'fullPath' => private_group.full_path }, fields_private_group) },
+          { query: graphql_query_for('group', { 'fullPath' => public_group.full_path }, fields_public_group) }
+        ]
+      end
+
+      let_it_be(:public_project) { create(:project, group: public_group) }
+      let_it_be(:private_project) { create(:project, group: private_group) }
+
+      let_it_be(:vulnerability_1) { create(:vulnerability, :dismissed, :critical_severity, :with_notes, notes_count: 2, project: public_project) }
+      let_it_be(:vulnerability_2) { create(:vulnerability, :confirmed, :high_severity, :with_notes, notes_count: 3, project: public_project) }
+      let_it_be(:vulnerability_3) { create(:vulnerability, :dismissed, :medium_severity, :with_notes, notes_count: 4, project: private_project) }
+      let_it_be(:vulnerability_4) { create(:vulnerability, :confirmed, :low_severity, :with_notes, notes_count: 7, project: private_project) }
+
+      let_it_be(:vulnerability_statistic_1) { create(:vulnerability_statistic, :grade_c, project: public_project) }
+      let_it_be(:vulnerability_statistic_2) { create(:vulnerability_statistic, :grade_d, project: private_project) }
+
+      let(:first_graphql_data) do
+        json_response.first['data']
+      end
+
+      let(:second_graphql_data) do
+        json_response.last['data']
+      end
+
+      let(:expected_private_group_response) do
+        [
+          {
+            'count' => 1,
+            'grade' => 'D',
+            'projects' => {
+              'nodes' => [
+                {
+                  'confirmedVulnerabilities' => {
+                    'nodes' => [
+                      { 'id' => vulnerability_4.to_global_id.to_s, 'userNotesCount' => 7 }
+                    ]
+                  },
+                  'dismissedVulnerabilities' => {
+                    'nodes' => [
+                      { 'id' => vulnerability_3.to_global_id.to_s, 'userNotesCount' => 4 }
+                    ]
+                  }
+                }
+              ]
+            }
+          },
+          {
+            'count' => 0,
+            'grade' => 'C',
+            'projects' => {
+              'nodes' => []
+            }
+          }
+        ]
+      end
+
+      let(:expected_public_group_response) do
+        [
+          {
+            'count' => 0,
+            'grade' => 'D',
+            'projects' => {
+              'nodes' => []
+            }
+          },
+          {
+            'count' => 1,
+            'grade' => 'C',
+            'projects' => {
+              'nodes' => [
+                {
+                  'vulnerabilities' => {
+                    'nodes' => [
+                      { 'id' => vulnerability_1.to_global_id.to_s, 'userNotesCount' => 2 },
+                      { 'id' => vulnerability_2.to_global_id.to_s, 'userNotesCount' => 3 }
+                    ]
+                  }
+                }
+              ]
+            }
+          }
+        ]
+      end
+
+      before do
+        public_group.add_developer(user)
+        private_group.add_developer(user)
+        stub_licensed_features(security_dashboard: true)
+
+        post_multiplex(queries, current_user: user)
+      end
+
+      it 'finds vulnerability grades for only projects that were added to instance security dashboard', :aggregate_failures do
+        expect(first_graphql_data.dig('group', 'vulnerabilityGrades')).to match_array(expected_private_group_response)
+        expect(second_graphql_data.dig('group', 'vulnerabilityGrades')).to match_array(expected_public_group_response)
+      end
+
+      it 'returns a successful response', :aggregate_failures do
+        expect(response).to have_gitlab_http_status(:success)
+        expect(graphql_errors).to eq([nil, nil])
+      end
+    end
   end
 end

ee/spec/requests/api/graphql/instance_security_dashboard_spec.rb

@@ -42,6 +42,147 @@ RSpec.describe 'Query.instanceSecurityDashboard.projects' do
         expect(projects).to eq([{ "id" => GitlabSchema.id_from_object(project).to_s }])
       end
     end
+
+    context 'when loading vulnerabilityGrades alongside with Vulnerability.userNotesCount' do
+      let(:fields) do
+        <<~QUERY
+        allGrades: vulnerabilityGrades {
+          grade
+          count
+          projects {
+            nodes {
+              vulnerabilities {
+                nodes {
+                  id
+                  userNotesCount
+                }
+              }
+            }
+          }
+        }
+        withVulnerabilitiesByState: vulnerabilityGrades {
+          grade
+          count
+          projects {
+            nodes {
+              confirmedVulnerabilities: vulnerabilities(state: CONFIRMED) {
+                nodes {
+                  id
+                  userNotesCount
+                }
+              }
+              dismissedVulnerabilities: vulnerabilities(state: DISMISSED) {
+                nodes {
+                  id
+                  userNotesCount
+                }
+              }
+            }
+          }
+        }
+        QUERY
+      end
+
+      let(:query) do
+        graphql_query_for('instanceSecurityDashboard', nil, fields)
+      end
+
+      let_it_be(:vulnerability_1) { create(:vulnerability, :dismissed, :critical_severity, :with_notes, notes_count: 2, project: project) }
+      let_it_be(:vulnerability_2) { create(:vulnerability, :confirmed, :high_severity, :with_notes, notes_count: 3, project: project) }
+      let_it_be(:vulnerability_3) { create(:vulnerability, :confirmed, :medium_severity, :with_notes, notes_count: 7, project: other_project) }
+
+      let_it_be(:vulnerability_statistic_1) { create(:vulnerability_statistic, :grade_c, project: project) }
+      let_it_be(:vulnerability_statistic_2) { create(:vulnerability_statistic, :grade_d, project: other_project) }
+
+      it_behaves_like 'a working graphql query' do
+        let(:expected_response) do
+          {
+            'allGrades' => [
+              {
+                'count' => 1,
+                'grade' => 'C',
+                'projects' => {
+                  'nodes' => [
+                    {
+                      'vulnerabilities' => {
+                        'nodes' => [
+                          { 'id' => vulnerability_1.to_global_id.to_s, 'userNotesCount' => 2 },
+                          { 'id' => vulnerability_2.to_global_id.to_s, 'userNotesCount' => 3 }
+                        ]
+                      }
+                    }
+                  ]
+                }
+              },
+              {
+                'count' => 1,
+                'grade' => 'D',
+                'projects' => {
+                  'nodes' => [
+                    {
+                      'vulnerabilities' => {
+                        'nodes' => [
+                          { 'id' => vulnerability_3.to_global_id.to_s, 'userNotesCount' => 7 }
+                        ]
+                      }
+                    }
+                  ]
+                }
+              }
+            ],
+            'withVulnerabilitiesByState' => [
+              {
+                'count' => 1,
+                'grade' => 'C',
+                'projects' => {
+                  'nodes' => [
+                    {
+                      'confirmedVulnerabilities' => {
+                        'nodes' => [
+                          { 'id' => vulnerability_2.to_global_id.to_s, 'userNotesCount' => 3 }
+                        ]
+                      },
+                      'dismissedVulnerabilities' => {
+                        'nodes' => [
+                          { 'id' => vulnerability_1.to_global_id.to_s, 'userNotesCount' => 2 }
+                        ]
+                      }
+                    }
+                  ]
+                }
+              },
+              {
+                'count' => 1,
+                'grade' => 'D',
+                'projects' => {
+                  'nodes' => [
+                    {
+                      'confirmedVulnerabilities' => {
+                        'nodes' => [
+                          { 'id' => vulnerability_3.to_global_id.to_s, 'userNotesCount' => 7 }
+                        ]
+                      },
+                      'dismissedVulnerabilities' => { 'nodes' => [] }
+                    }
+                  ]
+                }
+              }
+            ]
+          }
+        end
+
+        before do
+          user.security_dashboard_projects << other_project
+
+          post_graphql(query, current_user: current_user)
+        end
+
+        it 'finds vulnerability grades for only projects that were added to instance security dashboard', :aggregate_failures do
+          expect(graphql_data.dig('instanceSecurityDashboard', 'allGrades')).to match_array(expected_response['allGrades'])
+          expect(graphql_data.dig('instanceSecurityDashboard', 'withVulnerabilitiesByState')).to match_array(expected_response['withVulnerabilitiesByState'])
+        end
+      end
+    end
   end
 
   context 'with no user' do