Extra lazy calculation

-Don't even add up the direct sums
until we're ready to calculate everything
- Modify tests
- Make constants consistent COUNT_FACET is just COUNT
- Add limit to db call

doc/api/graphql/reference/gitlab_schema.graphql

@@ -2023,11 +2023,6 @@ type Epic implements Noteable {
   """
   hasIssues: Boolean!
 
-  """
-  Current health status. Available only when feature flag `save_issuable_health_status` is enabled.
-  """
-  healthStatus: HealthStatus
-
   """
   ID of the epic
   """

doc/api/graphql/reference/index.md

@@ -327,7 +327,6 @@ Represents an epic.
 | `group` | Group! | Group to which the epic belongs |
 | `hasChildren` | Boolean! | Indicates if the epic has children |
 | `hasIssues` | Boolean! | Indicates if the epic has direct issues |
-| `healthStatus` | HealthStatus | Current health status. Available only when feature flag `save_issuable_health_status` is enabled. |
 | `id` | ID! | ID of the epic |
 | `iid` | ID! | Internal ID of the epic |
 | `parent` | Epic | Parent epic of the epic |

ee/lib/gitlab/graphql/aggregations/epics/epic_node.rb

@@ -11,141 +11,84 @@ module Gitlab
           include Gitlab::Utils::StrongMemoize
 
           attr_reader :epic_id, :epic_state_id, :epic_info_flat_list, :parent_id,
-                      :direct_count_totals, :direct_weight_sum_totals, # only counts/weights of direct issues and child epic counts
                       :count_aggregate, :weight_sum_aggregate
 
           attr_accessor :children, :calculated_count_totals, :calculated_weight_sum_totals
 
           def initialize(epic_id, flat_info_list)
             # epic aggregate records from the DB loader look like the following:
-            # { 1 => [{iid: 1, epic_state_id: 1, issues_count: 1, issues_weight_sum: 2, parent_id: nil, state_id: 2}] ... }
+            # { 1 => [{iid: 1, parent_id: nil, epic_state_id: 1, issues_count: 1, issues_weight_sum: 2, issues_state_id: 2}] ... }
             # They include the sum of each epic's direct issues, grouped by status,
             # so in order to get a sum of the entire tree, we have to add that up recursively
             @epic_id = epic_id
             @epic_info_flat_list = flat_info_list
             @children = []
-            @direct_count_totals = []
-            @direct_weight_sum_totals = []
+            @sums = {}
 
             set_epic_attributes(flat_info_list.first) # there will always be one
           end
 
-          def assemble_issue_totals
-            # this is a representation of the epic's
-            # direct child issues and epics that have come from the DB
-            [OPENED_ISSUE_STATE, CLOSED_ISSUE_STATE].each do |issue_state|
-              matching_issue_state_entry = epic_info_flat_list.find do |epic_info_node|
-                epic_info_node[:issues_state_id] == issue_state
-              end || {}
-
-              create_sum_if_needed(WEIGHT_SUM, issue_state, ISSUE_TYPE, matching_issue_state_entry.fetch(:issues_weight_sum, 0))
-              create_sum_if_needed(COUNT, issue_state, ISSUE_TYPE, matching_issue_state_entry.fetch(:issues_count, 0))
-            end
-          end
-
-          def assemble_epic_totals
-            [OPENED_EPIC_STATE, CLOSED_EPIC_STATE].each do |epic_state|
-              create_sum_if_needed(COUNT, epic_state, EPIC_TYPE, children.select { |node| node.epic_state_id == epic_state }.count)
-            end
-          end
-
-          def aggregate_count(tree)
+          def aggregate_count
             strong_memoize(:count_aggregate) do
-              calculate_recursive_sums(COUNT, tree)
               OpenStruct.new({
                 opened_issues: sum_objects(COUNT, OPENED_ISSUE_STATE, ISSUE_TYPE),
                 closed_issues: sum_objects(COUNT, CLOSED_ISSUE_STATE, ISSUE_TYPE),
                 opened_epics: sum_objects(COUNT, OPENED_EPIC_STATE, EPIC_TYPE),
                 closed_epics: sum_objects(COUNT, CLOSED_EPIC_STATE, EPIC_TYPE)
-              })
+             })
             end
           end
 
-          def aggregate_weight_sum(tree)
+          def aggregate_weight_sum
             strong_memoize(:weight_sum_aggregate) do
-              calculate_recursive_sums(WEIGHT_SUM, tree)
               OpenStruct.new({
                 opened_issues: sum_objects(WEIGHT_SUM, OPENED_ISSUE_STATE, ISSUE_TYPE),
                 closed_issues: sum_objects(WEIGHT_SUM, CLOSED_ISSUE_STATE, ISSUE_TYPE)
-              })
+               })
             end
           end
 
-          def direct_totals(facet)
-            # Sums of only child issues and immediate child epics (but not their issues
-            # )
-            strong_memoize(:"direct_#{facet}_totals") do
-              []
-            end
-          end
-
-          def calculated_totals(facet)
-            if facet == COUNT
-              return calculated_count_totals
-            end
-
-            calculated_weight_sum_totals
-          end
-
-          def calculate_recursive_sums(facet, tree)
-            return calculated_totals(facet) if calculated_totals(facet)
-
-            sum_total = []
-            children.each do |child|
-              child_sums = child.calculate_recursive_sums(facet, tree)
-              sum_total.concat(child_sums)
-            end
-            sum_total.concat(direct_totals(facet))
-            set_calculated_total(facet, sum_total)
-          end
-
-          def inspect
+          def to_s
             {
                 epic_id: @epic_id,
                 parent_id: @parent_id,
-                direct_count_totals: direct_count_totals,
-                direct_weight_sum_totals: direct_weight_sum_totals,
                 children: children,
                 object_id: object_id
-            }.to_json
+            }.to_s
           end
 
-          alias_method :to_s, :inspect
-
-          private
-
           def sum_objects(facet, state, type)
-            sums = calculated_totals(facet) || []
-            return 0 if sums.empty?
+            key = [facet, state, type]
+            return @sums[key] if @sums[key]
 
-            sums.inject(0) do |result, sum|
-              result += sum.value if sum.state == state && sum.type == type
-              result
+            direct_sum = value_from_records(*key)
+            sum_from_children = children.inject(0) do |total, child|
+              total += child.sum_objects(*key)
+              total
             end
-          end
-
-          def create_sum_if_needed(facet, state, type, value)
-            return if value.nil? || value < 1
 
-            direct_totals(facet) << Sum.new(facet, state, type, value)
+            @sums[key] = direct_sum + sum_from_children
           end
 
+          private
+
           def set_epic_attributes(record)
             @epic_state_id = record[:epic_state_id]
             @parent_id = record[:parent_id]
           end
 
-          def set_calculated_total(facet, calculated_sums)
-            if facet == COUNT
-              @calculated_count_totals = calculated_sums
+          def value_from_records(facet, state, type)
+            # DB records look like:
+            # {iid: 1, epic_state_id: 1, issues_count: 1, issues_weight_sum: 2, parent_id: nil, issues_state_id: 2}
+            if type == EPIC_TYPE
+              # can only be COUNT
+              children.select { |node| node.epic_state_id == state }.count
             else
-              @calculated_weight_sum_totals = calculated_sums
-            end
-          end
+              matching_record = epic_info_flat_list.find do |record|
+                record[:issues_state_id] == state
+              end || {}
 
-          Sum = Struct.new(:facet, :state, :type, :value) do
-            def inspect
-              "<Sum facet=#{facet}, state=#{state}, type=#{type}, value=#{value}>"
+              matching_record.fetch("issues_#{facet}".to_sym, 0)
             end
           end
         end

ee/lib/gitlab/graphql/aggregations/epics/lazy_epic_aggregate.rb

@@ -67,43 +67,34 @@ module Gitlab
 
             # Fire off the db query and get the results (grouped by epic_id and facet)
             raw_epic_aggregates = Gitlab::Graphql::Loaders::BulkEpicAggregateLoader.new(epic_ids: pending_ids).execute
-
-            # Assemble the tree and sum immediate child epic/issues
-            create_structure_from(raw_epic_aggregates)
-
+            create_epic_nodes(raw_epic_aggregates)
             @lazy_state[:pending_ids].clear
           end
 
-          def create_structure_from(aggregate_records)
-            # create EpicNode object for each epic id
+          def create_epic_nodes(aggregate_records)
             aggregate_records.each do |epic_id, aggregates|
               next if aggregates.blank?
 
               tree[epic_id] = EpicNode.new(epic_id, aggregates)
             end
 
-            assemble_direct_child_totals
+            relate_parents_and_children
           end
 
-          def assemble_direct_child_totals
+          def relate_parents_and_children
             tree.each do |_, node|
               parent = tree[node.parent_id]
               next if parent.nil?
 
               parent.children << node
             end
-
-            tree.each do |_, node|
-              node.assemble_epic_totals
-              node.assemble_issue_totals
-            end
           end
 
           def aggregate_object(node)
             if @facet == COUNT
-              node.aggregate_count(tree)
+              node.aggregate_count
             else
-              node.aggregate_weight_sum(tree)
+              node.aggregate_weight_sum
             end
           end
         end

ee/lib/gitlab/graphql/loaders/bulk_epic_aggregate_loader.rb

@@ -6,6 +6,8 @@ module Gitlab
       class BulkEpicAggregateLoader
         include ::Gitlab::Graphql::Aggregations::Epics::Constants
 
+        MAXIMUM_LOADABLE = 100_001
+
         attr_reader :target_epic_ids, :results
 
         # This class retrieves each epic and its child epics recursively
@@ -25,8 +27,12 @@ module Gitlab
             .left_joins(epic_issues: :issue)
             .group("issues.state_id", "epics.id", "epics.iid", "epics.parent_id", "epics.state_id")
             .select("epics.id, epics.iid, epics.parent_id, epics.state_id AS epic_state_id, issues.state_id AS issues_state_id, COUNT(issues) AS issues_count, SUM(COALESCE(issues.weight, 0)) AS issues_weight_sum")
+            .limit(MAXIMUM_LOADABLE)
+
+          raw_results = raw_results.map { |record| record.attributes.with_indifferent_access }
+
+          raise ArgumentError.new("There are too many records to load. Please select fewer epics or contact your administrator.") if raw_results.count == MAXIMUM_LOADABLE
 
-          raw_results = raw_results.map(&:attributes).map(&:with_indifferent_access)
           @results = raw_results.group_by { |record| record[:id] }
         end
         # rubocop: enable CodeReuse/ActiveRecord

ee/spec/lib/gitlab/graphql/aggregations/epics/epic_node_spec.rb

@@ -32,194 +32,105 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
     it_behaves_like 'setting attributes based on the first record', { epic_state_id: CLOSED_EPIC_STATE, parent_id: 2 }
   end
 
-  describe '#assemble_issue_totals' do
-    subject { described_class.new(epic_id, fake_data) }
-
-    context 'an epic with no issues' do
-      let(:fake_data) do
-        [
-            { iid: epic_iid, epic_state_id: OPENED_EPIC_STATE, issues_count: 0, issues_weight_sum: 0, parent_id: nil, issues_state_id: nil }
-        ]
-      end
-
-      it 'does not create any totals' do
-        subject.assemble_issue_totals
+  describe 'recursive totals' do
+    subject { described_class.new(epic_id, [{ parent_id: nil, epic_state_id: CLOSED_EPIC_STATE }]) }
 
-        expect(subject.direct_totals(COUNT_FACET).count).to eq 0
-        expect(subject.direct_totals(WEIGHT_SUM_FACET).count).to eq 0
-      end
+    before do
+      allow(subject).to receive(:epic_info_flat_list).and_return(flat_info)
     end
 
-    context 'an epic with issues' do
-      context 'with a nonzero count but a zero weight' do
-        let(:fake_data) do
-          [
-              { iid: epic_iid, epic_state_id: OPENED_EPIC_STATE, issues_count: 1, issues_weight_sum: 0, parent_id: nil, issues_state_id: OPENED_ISSUE_STATE }
-          ]
-        end
+    context 'an epic with no child epics' do
+      context 'with no child issues', :aggregate_results do
+        let(:flat_info) { [] }
 
-        it 'creates no sums for the weight if the issues have 0 weight' do
-          subject.assemble_issue_totals
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
 
-          expect(subject.direct_totals(COUNT_FACET).count).to eq 1
-          expect(subject.direct_totals(WEIGHT_SUM_FACET).count).to eq 0
-          expect(subject).to have_direct_total(ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 1)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
       end
 
-      context 'with a nonzero count and nonzero weight for a single state' do
-        let(:fake_data) do
+      context 'with an issue with 0 weight', :aggregate_results do
+        let(:flat_info) do
           [
-              { iid: epic_iid, epic_state_id: OPENED_EPIC_STATE, issues_count: 1, issues_weight_sum: 2, parent_id: nil, issues_state_id: OPENED_ISSUE_STATE }
+              record_for(epic_id: epic_id, parent_id: nil, epic_state_id: CLOSED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 0)
           ]
         end
 
-        it 'creates two sums' do
-          subject.assemble_issue_totals
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 1)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
 
-          expect(subject).to have_direct_total(ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 1)
-          expect(subject).to have_direct_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 2)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
       end
 
-      context 'with a nonzero count and nonzero weight for multiple states' do
-        let(:fake_data) do
+      context 'with an issue with nonzero weight' do
+        let(:flat_info) do
           [
-              { iid: epic_iid, epic_state_id: OPENED_EPIC_STATE, issues_count: 1, issues_weight_sum: 2, parent_id: nil, issues_state_id: OPENED_ISSUE_STATE },
-              { iid: epic_iid, epic_state_id: OPENED_EPIC_STATE, issues_count: 3, issues_weight_sum: 5, parent_id: nil, issues_state_id: CLOSED_ISSUE_STATE }
+              record_for(epic_id: epic_id, parent_id: nil, epic_state_id: CLOSED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 2)
           ]
         end
 
-        it 'creates two sums' do
-          subject.assemble_issue_totals
-
-          expect(subject).to have_direct_total(ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 1)
-          expect(subject).to have_direct_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 2)
-          expect(subject).to have_direct_total(ISSUE_TYPE, COUNT_FACET, CLOSED_ISSUE_STATE, 3)
-          expect(subject).to have_direct_total(ISSUE_TYPE, WEIGHT_SUM_FACET, CLOSED_ISSUE_STATE, 5)
-        end
-      end
-    end
-  end
-
-  describe '#assemble_epic_totals' do
-    subject { described_class.new(epic_id, [{ parent_id: nil, epic_state_id: CLOSED_EPIC_STATE }]) }
-
-    context 'with a child epic' do
-      let(:child_epic_id) { 45 }
-      let!(:child_epic_node) { described_class.new(child_epic_id, [{ parent_id: epic_id, epic_state_id: CLOSED_EPIC_STATE }]) }
-
-      before do
-        subject.children << child_epic_node
-      end
-
-      it 'adds up the number of the child epics' do
-        subject.assemble_epic_totals
-
-        expect(subject).to have_direct_total(EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
-      end
-    end
-  end
-
-  describe '#calculate_recursive_sums' do
-    subject { described_class.new(epic_id, [{ parent_id: nil, epic_state_id: CLOSED_EPIC_STATE }]) }
-
-    before do
-      allow(subject).to receive(:direct_totals).with(COUNT_FACET).and_return(immediate_count_totals)
-      allow(subject).to receive(:direct_totals).with(WEIGHT_SUM_FACET).and_return(immediate_weight_sum_totals)
-    end
-
-    shared_examples 'returns calculated totals by facet' do |facet, count|
-      it 'returns a calculated_count_total' do
-        subject.calculate_recursive_sums(facet, tree)
-
-        expect(subject.calculated_totals(facet).count).to eq count
-      end
-    end
-
-    context 'an epic with no child epics' do
-      let(:tree) do
-        { epic_id => subject }
-      end
-
-      context 'with no child issues' do
-        let(:immediate_count_totals) { [] }
-        let(:immediate_weight_sum_totals) { [] }
-
-        it_behaves_like 'returns calculated totals by facet', COUNT_FACET, 0
-        it_behaves_like 'returns calculated totals by facet', WEIGHT_SUM_FACET, 0
-      end
-
-      context 'with an issue with 0 weight' do
-        let(:immediate_count_totals) do
-          [Gitlab::Graphql::Aggregations::Epics::EpicNode::Sum.new(COUNT_FACET, CLOSED_EPIC_STATE, ISSUE_TYPE, 1)]
-        end
-        let(:immediate_weight_sum_totals) { [] }
-
-        it_behaves_like 'returns calculated totals by facet', COUNT_FACET, 1
-        it_behaves_like 'returns calculated totals by facet', WEIGHT_SUM_FACET, 0
-      end
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 1)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 2)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
 
-      context 'with an issue with nonzero weight' do
-        let(:immediate_count_totals) do
-          [
-            Gitlab::Graphql::Aggregations::Epics::EpicNode::Sum.new(COUNT_FACET, CLOSED_EPIC_STATE, ISSUE_TYPE, 1)
-          ]
-        end
-        let(:immediate_weight_sum_totals) do
-          [
-            Gitlab::Graphql::Aggregations::Epics::EpicNode::Sum.new(WEIGHT_SUM_FACET, CLOSED_EPIC_STATE, ISSUE_TYPE, 2)
-          ]
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
-
-        it_behaves_like 'returns calculated totals by facet', COUNT_FACET, 1
-        it_behaves_like 'returns calculated totals by facet', WEIGHT_SUM_FACET, 1
       end
     end
 
     context 'an epic with child epics' do
       let(:child_epic_id) { 45 }
-      let(:tree) do
-        { epic_id => subject, child_epic_id => child_epic_node }
-      end
-      let(:child_epic_node) { described_class.new(child_epic_id, [{ parent_id: epic_id, epic_state_id: CLOSED_EPIC_STATE }]) }
-      let(:immediate_count_totals) do
-        [ # only one opened epic, the child
-          Gitlab::Graphql::Aggregations::Epics::EpicNode::Sum.new(COUNT_FACET, OPENED_EPIC_STATE, EPIC_TYPE, 1)
+      let(:child_epic_node) { described_class.new(child_epic_id, child_flat_info) }
+      let(:flat_info) do
+        [
+            record_for(epic_id: epic_id, parent_id: nil, epic_state_id: OPENED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 0, issues_weight_sum: 0)
         ]
       end
-      let(:immediate_weight_sum_totals) { [] }
 
       before do
         subject.children << child_epic_node
-        allow(child_epic_node).to receive(:direct_totals).with(COUNT_FACET).and_return(child_count_totals)
-        allow(child_epic_node).to receive(:direct_totals).with(WEIGHT_SUM_FACET).and_return(child_weight_sum_totals)
       end
 
       context 'with a child that has issues of nonzero weight' do
-        let(:child_count_totals) do
+        let(:child_flat_info) do
           [
-            Gitlab::Graphql::Aggregations::Epics::EpicNode::Sum.new(COUNT_FACET, OPENED_ISSUE_STATE, ISSUE_TYPE, 1)
+            record_for(epic_id: epic_id, parent_id: nil, epic_state_id: OPENED_EPIC_STATE, issues_state_id: OPENED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 2)
           ]
         end
-        let(:child_weight_sum_totals) do
-          [
-            Gitlab::Graphql::Aggregations::Epics::EpicNode::Sum.new(WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, ISSUE_TYPE, 2)
-          ]
-        end
-
-        it 'returns the correct count total' do
-          subject.calculate_recursive_sums(COUNT_FACET, tree)
-
-          expect(subject).to have_calculated_total(ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 1)
-        end
-
-        it 'returns the correct weight sum total' do
-          subject.calculate_recursive_sums(WEIGHT_SUM_FACET, tree)
 
-          expect(subject).to have_calculated_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 2)
+        it 'has the correct aggregates', :aggregate_failures do
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 1)
+          expect(subject).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 2)
+          expect(subject).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, OPENED_EPIC_STATE, 1)
+          expect(subject).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
       end
     end
   end
+
+  def record_for(epic_id:, parent_id:, epic_state_id:, issues_state_id:, issues_count:, issues_weight_sum:)
+    {
+      epic_id: epic_id,
+      issues_count: issues_count,
+      issues_weight_sum: issues_weight_sum,
+      parent_id: parent_id,
+      issues_state_id: issues_state_id,
+      epic_state_id: epic_state_id
+    }
+  end
 end

ee/spec/lib/gitlab/graphql/aggregations/epics/lazy_epic_aggregate_spec.rb

@@ -21,20 +21,20 @@ describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
 
     context 'with valid facets :weight_sum or :count' do
       specify 'as a symbol', :aggregate_failures do
-        [:weight_sum, :count].each do |valid_facet|
-          described_class.new(query_ctx, epic_id, valid_facet)
+        [WEIGHT_SUM, COUNT].each do |valid_facet|
+          expect { described_class.new(query_ctx, epic_id, valid_facet) }.not_to raise_error
         end
       end
 
       specify 'as a string', :aggregate_failures do
         %w(weight_sum count).each do |valid_facet|
-          described_class.new(query_ctx, epic_id, valid_facet)
+          expect { described_class.new(query_ctx, epic_id, valid_facet) }.not_to raise_error
         end
       end
     end
 
     it 'adds the epic_id to lazy state' do
-      described_class.new(query_ctx, epic_id, :count)
+      described_class.new(query_ctx, epic_id, COUNT)
 
       expect(query_ctx[:lazy_epic_aggregate][:pending_ids]).to match [epic_id]
     end
@@ -46,7 +46,7 @@ describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
     end
     let(:epic_info_node) { Gitlab::Graphql::Aggregations::Epics::EpicNode.new(epic_id, [single_record] ) }
 
-    subject { described_class.new(query_ctx, epic_id, :count) }
+    subject { described_class.new(query_ctx, epic_id, COUNT) }
 
     before do
       subject.instance_variable_set(:@lazy_state, fake_state)
@@ -79,8 +79,9 @@ describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
       end
 
       before do
-        allow(Gitlab::Graphql::Aggregations::Epics::EpicNode).to receive(:aggregate_count).and_call_original
-        expect_any_instance_of(Gitlab::Graphql::Loaders::BulkEpicAggregateLoader).to receive(:execute).and_return(fake_data)
+        expect_next_instance_of(Gitlab::Graphql::Loaders::BulkEpicAggregateLoader) do |loader|
+          expect(loader).to receive(:execute).and_return(fake_data)
+        end
       end
 
       it 'clears the pending IDs' do
@@ -94,53 +95,38 @@ describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
       it 'creates the parent-child associations', :aggregate_failures do
         subject.epic_aggregate
 
-        lazy_state = subject.instance_variable_get(:@lazy_state)
-        tree = lazy_state[:tree]
-
         expect(tree[child_epic_id].parent_id).to eq epic_id
         expect(tree[epic_id].children.map(&:epic_id)).to match_array([child_epic_id])
       end
 
       context 'for a parent-child relationship' do
-        it 'assembles direct sums', :aggregate_failures do
-          subject.epic_aggregate
-
-          lazy_state = subject.instance_variable_get(:@lazy_state)
-          tree = lazy_state[:tree]
-
-          expect(tree[epic_id]).to have_direct_total(EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
-          expect(tree[epic_id]).to have_direct_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 5)
-          expect(tree[epic_id]).to have_direct_total(EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
-
-          expect(tree[child_epic_id]).to have_direct_total(ISSUE_TYPE, COUNT_FACET, CLOSED_ISSUE_STATE, 4)
-          expect(tree[child_epic_id]).to have_direct_total(ISSUE_TYPE, WEIGHT_SUM_FACET, CLOSED_ISSUE_STATE, 17)
-        end
-
         it 'assembles recursive sums for the parent', :aggregate_failures do
           subject.epic_aggregate
 
-          lazy_state = subject.instance_variable_get(:@lazy_state)
-          tree = lazy_state[:tree]
-
-          expect(tree[epic_id]).to have_aggregate(tree, ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 2)
-          expect(tree[epic_id]).to have_aggregate(tree, ISSUE_TYPE, COUNT_FACET, CLOSED_ISSUE_STATE, 4)
-          expect(tree[epic_id]).to have_aggregate(tree, ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 5)
-          expect(tree[epic_id]).to have_aggregate(tree, ISSUE_TYPE, WEIGHT_SUM_FACET, CLOSED_ISSUE_STATE, 17)
-          expect(tree[epic_id]).to have_aggregate(tree, EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 2)
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 4)
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 5)
+          expect(tree[epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 17)
+          expect(tree[epic_id]).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 1)
         end
       end
 
       context 'for a standalone epic with no issues' do
-        it 'assembles direct totals', :aggregate_failures do
+        it 'assembles recursive sums', :aggregate_failures do
           subject.epic_aggregate
 
-          lazy_state = subject.instance_variable_get(:@lazy_state)
-          tree = lazy_state[:tree]
-
-          expect(tree[other_epic_id].direct_count_totals).to be_empty
-          expect(tree[other_epic_id].direct_weight_sum_totals).to be_empty
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, OPENED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, COUNT, CLOSED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, OPENED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(ISSUE_TYPE, WEIGHT_SUM, CLOSED_ISSUE_STATE, 0)
+          expect(tree[other_epic_id]).to have_aggregate(EPIC_TYPE, COUNT, CLOSED_EPIC_STATE, 0)
         end
       end
     end
   end
+
+  def tree
+    lazy_state = subject.instance_variable_get(:@lazy_state)
+    lazy_state[:tree]
+  end
 end

ee/spec/lib/gitlab/graphql/loaders/bulk_epic_aggregate_loader_spec.rb

@@ -87,6 +87,12 @@ describe Gitlab::Graphql::Loaders::BulkEpicAggregateLoader do
       expect(result.keys).to match_array([parent_epic.id, epic_with_issues.id, epic_without_issues.id])
     end
 
+    it 'errors when the number of retrieved records exceeds the maximum' do
+      stub_const("Gitlab::Graphql::Loaders::BulkEpicAggregateLoader::MAXIMUM_LOADABLE", 1)
+
+      expect { subject.execute }.to raise_error(ArgumentError, /too many records/)
+    end
+
     context 'testing for a single database query' do
       it 'does not repeat database queries for subepics' do
         recorder = ActiveRecord::QueryRecorder.new { described_class.new(epic_ids: epic_with_issues.id).execute }
@@ -128,6 +134,14 @@ describe Gitlab::Graphql::Loaders::BulkEpicAggregateLoader do
   end
 
   def result_for(epic, issues_state:, issues_count:, issues_weight_sum:)
-    { id: epic.id, iid: epic.iid, issues_count: issues_count, issues_weight_sum: issues_weight_sum, parent_id: epic.parent_id, issues_state_id: issues_state, epic_state_id: Epic.available_states[epic.state_id] }.stringify_keys
+    {
+      id: epic.id,
+      iid: epic.iid,
+      issues_count: issues_count,
+      issues_weight_sum: issues_weight_sum,
+      parent_id: epic.parent_id,
+      issues_state_id: issues_state,
+      epic_state_id: Epic.available_states[epic.state_id]
+    }.stringify_keys
   end
 end

ee/spec/support/matchers/ee/epic_aggregate_matchers.rb

 # frozen_string_literal: true
 
-%w[direct calculated].each do |total_type|
-  RSpec::Matchers.define :"have_#{total_type}_total" do |type, facet, state, expected_value|
-    match do |epic_node_result|
-      expect(epic_node_result).not_to be_nil
-      totals = epic_node_result.public_send("#{total_type}_totals", facet)
-      expect(totals).not_to be_empty
-
-      matching = totals.select { |sum| sum.type == type && sum.facet == facet && sum.state == state && sum.value == expected_value }
-      expect(matching).not_to be_empty
-    end
-
-    failure_message do |epic_node_result|
-      if epic_node_result.nil?
-        "expected for there to be an epic node, but it is nil"
-      else
-        totals = epic_node_result.public_send("#{total_type}_totals", facet)
-        <<~FAILURE_MSG
-          expected epic node with id #{epic_node_result.epic_id} to have a sum with facet '#{facet}', state '#{state}', type '#{type}' and value '#{expected_value}'. Has #{totals.count} #{total_type} sum objects#{", none of which match" if totals.count > 0}.
-          Sums: #{totals.inspect}
-        FAILURE_MSG
-      end
-    end
-  end
-end
-
-RSpec::Matchers.define :have_aggregate do |tree, type, facet, state, expected_value|
+RSpec::Matchers.define :have_aggregate do |type, facet, state, expected_value|
   match do |epic_node_result|
-    aggregate_object = epic_node_result.public_send(:"aggregate_#{facet}", tree)
+    aggregate_object = epic_node_result.public_send(:"aggregate_#{facet}")
     expect(aggregate_object.public_send(method_name(type, state))).to eq expected_value
   end
 
   failure_message do |epic_node_result|
-    aggregate_object = epic_node_result.public_send(:"aggregate_#{facet}", tree)
+    aggregate_object = epic_node_result.public_send(:"aggregate_#{facet}")
     aggregate_method = method_name(type, state)
     "Epic node with id #{epic_node_result.epic_id} called #{aggregate_method} on aggregate object. Value was expected to be #{expected_value} but was #{aggregate_object.send(aggregate_method)}."
   end

ee/spec/support/shared_contexts/epic_aggregate_constants.rb

@@ -9,6 +9,6 @@ shared_context 'includes EpicAggregate constants' do
   OPENED_ISSUE_STATE = Gitlab::Graphql::Aggregations::Epics::Constants::OPENED_ISSUE_STATE
   CLOSED_ISSUE_STATE = Gitlab::Graphql::Aggregations::Epics::Constants::CLOSED_ISSUE_STATE
 
-  WEIGHT_SUM_FACET = Gitlab::Graphql::Aggregations::Epics::Constants::WEIGHT_SUM
-  COUNT_FACET = Gitlab::Graphql::Aggregations::Epics::Constants::COUNT
+  WEIGHT_SUM = Gitlab::Graphql::Aggregations::Epics::Constants::WEIGHT_SUM
+  COUNT = Gitlab::Graphql::Aggregations::Epics::Constants::COUNT
 end