Apply reviewer comments

- Rename immediate to direct
- Another matcher for calculated as well as
 direct sums
- Remove unneeded constant methods
- Test epic node state
- Simplify epic tree load methods
- Combine epic children and direct total
assembly

doc/api/graphql/reference/gitlab_schema.graphql

@@ -1948,8 +1948,8 @@ type Epic implements Noteable {
   descendantCounts: EpicDescendantCount
 
   """
-  Total weight of open and closed descendant epic's issues. Available only when
-  feature flag `unfiltered_epic_aggregates` is enabled.
+  Total weight of open and closed issues in the epic and its descendants.
+  Available only when feature flag `unfiltered_epic_aggregates` is enabled.
   """
   descendantWeightSum: EpicDescendantWeights
 
@@ -2023,6 +2023,11 @@ 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

@@ -317,7 +317,7 @@ Represents an epic.
 | `closedAt` | Time | Timestamp of the epic's closure |
 | `createdAt` | Time | Timestamp of the epic's creation |
 | `descendantCounts` | EpicDescendantCount | Number of open and closed descendant epics and issues |
-| `descendantWeightSum` | EpicDescendantWeights | Total weight of open and closed descendant epic's issues. Available only when feature flag `unfiltered_epic_aggregates` is enabled. |
+| `descendantWeightSum` | EpicDescendantWeights | Total weight of open and closed issues in the epic and its descendants. Available only when feature flag `unfiltered_epic_aggregates` is enabled. |
 | `description` | String | Description of the epic |
 | `downvotes` | Int! | Number of downvotes the epic has received |
 | `dueDate` | Time | Due date of the epic |
@@ -327,6 +327,7 @@ 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

 # frozen_string_literal: true
 
-# This class represents an Epic's aggregate information (added up counts) about its child epics and immediate issues
+# This class represents an Epic's aggregate information (added up counts) about its child epics and direct issues
 
 module Gitlab
   module Graphql
@@ -11,7 +11,7 @@ module Gitlab
           include Gitlab::Utils::StrongMemoize
 
           attr_reader :epic_id, :epic_state_id, :epic_info_flat_list, :parent_id,
-                      :immediate_count_totals, :immediate_weight_sum_totals, # only counts/weights of immediate issues and child epic counts
+                      :direct_count_totals, :direct_weight_sum_totals, # only counts/weights of direct issues and child epic counts
                       :count_aggregate, :weight_sum_aggregate
 
           attr_accessor :child_ids, :calculated_count_totals, :calculated_weight_sum_totals
@@ -19,20 +19,20 @@ module Gitlab
           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}] ... }
-            # They include the sum of each epic's immediate issues, grouped by status,
+            # 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
             @child_ids = []
-            @immediate_count_totals = []
-            @immediate_weight_sum_totals = []
+            @direct_count_totals = []
+            @direct_weight_sum_totals = []
 
             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
-            # immediate child issues and epics that have come from the DB
+            # 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
@@ -71,8 +71,10 @@ module Gitlab
             end
           end
 
-          def immediate_totals(facet)
-            strong_memoize(:"immediate_#{facet}_totals") do
+          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
@@ -85,19 +87,6 @@ module Gitlab
             calculated_weight_sum_totals
           end
 
-          def to_s
-            {
-                epic_id: @epic_id,
-                parent_id: @parent_id,
-                immediate_count_totals: immediate_count_totals,
-                immediate_weight_sum_totals: immediate_weight_sum_totals,
-                child_ids: child_ids
-            }.to_json
-          end
-
-          alias_method :inspect, :to_s
-          alias_method :id, :epic_id
-
           def calculate_recursive_sums(facet, tree)
             return calculated_totals(facet) if calculated_totals(facet)
 
@@ -108,7 +97,7 @@ module Gitlab
               child_sums = child.calculate_recursive_sums(facet, tree)
               sum_total.concat(child_sums)
             end
-            sum_total.concat(immediate_totals(facet))
+            sum_total.concat(direct_totals(facet))
             set_calculated_total(facet, sum_total)
           end
 
@@ -116,16 +105,18 @@ module Gitlab
 
           def sum_objects(facet, state, type)
             sums = calculated_totals(facet) || []
-            matching = sums.select { |sum| sum.state == state && sum.type == type }
             return 0 if sums.empty?
 
-            matching.map(&:value).reduce(:+) || 0
+            sums.inject(0) do |result, sum|
+              result += sum.value if sum.state == state && sum.type == type
+              result
+            end
           end
 
           def create_sum_if_needed(facet, state, type, value)
             return if value.nil? || value < 1
 
-            immediate_totals(facet) << Sum.new(facet, state, type, value)
+            direct_totals(facet) << Sum.new(facet, state, type, value)
           end
 
           def set_epic_attributes(record)

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

@@ -9,13 +9,17 @@ module Gitlab
 
           attr_reader :facet, :epic_id, :lazy_state
 
+          PERMITTED_FACETS = [COUNT, WEIGHT_SUM].freeze
+
           # Because facets "count" and "weight_sum" share the same db query, but have a different graphql type object,
           # we can separate them and serve only the fields which are requested by the GraphQL query
           def initialize(query_ctx, epic_id, aggregate_facet)
             @epic_id = epic_id
 
-            raise ArgumentError.new("No aggregate facet provided. Please specify either #{COUNT} or #{WEIGHT_SUM}") unless aggregate_facet.present?
-            raise ArgumentError.new("Invalid aggregate facet #{aggregate_facet} provided. Please specify either #{COUNT} or #{WEIGHT_SUM}") unless [COUNT, WEIGHT_SUM].include?(aggregate_facet.to_sym)
+            error = validate_facet(aggregate_facet)
+            if error
+              raise ArgumentError.new("#{error}. Please specify either #{COUNT} or #{WEIGHT_SUM}")
+            end
 
             @facet = aggregate_facet.to_sym
 
@@ -33,19 +37,25 @@ module Gitlab
           def epic_aggregate
             # Check if the record was already loaded:
             # load from tree by epic
-            loaded_epic_info_node = @lazy_state[:tree][@epic_id]
-
-            if loaded_epic_info_node
-              # The pending IDs were already loaded,
-              # so return the result of that previous load
-              aggregate_object(loaded_epic_info_node)
-            else
+            unless tree[@epic_id]
               load_records_into_tree
             end
+
+            aggregate_object(tree[@epic_id])
           end
 
           private
 
+          def validate_facet(aggregate_facet)
+            unless aggregate_facet.present?
+              return "No aggregate facet provided."
+            end
+
+            unless PERMITTED_FACETS.include?(aggregate_facet.to_sym)
+              return "Invalid aggregate facet #{aggregate_facet} provided."
+            end
+          end
+
           def tree
             @lazy_state[:tree]
           end
@@ -62,38 +72,26 @@ module Gitlab
             create_structure_from(raw_epic_aggregates)
 
             @lazy_state[:pending_ids].clear
-
-            # Now, get the matching node and return its aggregate depending on the facet:
-            epic_node = @lazy_state[:tree][@epic_id]
-            aggregate_object(epic_node)
           end
 
           def create_structure_from(aggregate_records)
             # create EpicNode object for each epic id
             aggregate_records.each do |epic_id, aggregates|
-              next if aggregates.nil? || aggregates.empty?
+              next if aggregates.blank?
 
-              new_node = EpicNode.new(epic_id, aggregates)
-              tree[epic_id] = new_node
+              tree[epic_id] = EpicNode.new(epic_id, aggregates)
             end
 
-            associate_parents_and_children
-            assemble_immediate_totals
-          end
-
-          # each of the methods below are done one after the other
-          def associate_parents_and_children
-            tree.each do |epic_id, node|
-              node.child_ids = tree.select { |_, child_node| epic_id == child_node.parent_id }.keys
-            end
+            assemble_direct_child_totals
           end
 
-          def assemble_immediate_totals
+          def assemble_direct_child_totals
             tree.each do |_, node|
-              node.assemble_issue_totals
+              node_children = tree.select { |_, child_node| node.epic_id == child_node.parent_id }
+              node.child_ids = node_children.keys
+              node.assemble_epic_totals(node_children.values)
 
-              node_children = tree.select { |_, child_node| node.epic_id == child_node.parent_id }.values
-              node.assemble_epic_totals(node_children)
+              node.assemble_issue_totals
             end
           end
 

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

@@ -6,16 +6,18 @@ module Gitlab
       class BulkEpicAggregateLoader
         include ::Gitlab::Graphql::Aggregations::Epics::Constants
 
+        attr_reader :target_epic_ids, :results
+
         # This class retrieves each epic and its child epics recursively
         # It allows us to recreate the epic tree structure in POROs
         def initialize(epic_ids:)
           @results = {}
-          @target_epic_ids = epic_ids.present? ? [epic_ids].flatten.compact : []
+          @target_epic_ids = epic_ids
         end
 
         # rubocop: disable CodeReuse/ActiveRecord
         def execute
-          return {} unless target_epic_ids.any?
+          return {} unless target_epic_ids
 
           # We do a left outer join in order to capture epics with no issues
           # This is so we can aggregate the epic counts for every epic
@@ -25,27 +27,9 @@ module Gitlab
             .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")
 
           raw_results = raw_results.map(&:attributes).map(&:with_indifferent_access)
-          group_by_epic_id(raw_results)
-          @results
+          @results = raw_results.group_by { |record| record[:id] }
         end
         # rubocop: enable CodeReuse/ActiveRecord
-
-        private
-
-        attr_reader :target_epic_ids, :results
-
-        def group_by_epic_id(raw_records)
-          # for each id, populate with matching records
-          # change from a series of { id: x ... } to { x => [{...}, {...}] }
-          raw_records.map { |r| r[:id] }.uniq.each do |epic_id|
-            records = []
-            matching_records = raw_records.select { |record| record[:id] == epic_id }
-            matching_records.each do |record|
-              records << record.except(:id).to_h.with_indifferent_access
-            end
-            @results[epic_id] = records
-          end
-        end
       end
     end
   end

ee/spec/graphql/types/epic_type_spec.rb

@@ -11,7 +11,7 @@ describe GitlabSchema.types['Epic'] do
       closed_at created_at updated_at children has_children has_issues
       web_path web_url relation_path reference issues user_permissions
       notes discussions relative_position subscribed participants
-      descendant_counts descendant_weight_sum upvotes downvotes
+      descendant_counts descendant_weight_sum upvotes downvotes health_status
     ]
   end
 

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

@@ -24,6 +24,7 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
         new_node = described_class.new(epic_id, fake_data)
 
         expect(new_node.parent_id).to eq parent_id
+        expect(new_node.epic_state_id).to eq epic_state_id
       end
     end
 
@@ -44,8 +45,8 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
       it 'does not create any totals' do
         subject.assemble_issue_totals
 
-        expect(subject.immediate_totals(COUNT_FACET).count).to eq 0
-        expect(subject.immediate_totals(WEIGHT_SUM_FACET).count).to eq 0
+        expect(subject.direct_totals(COUNT_FACET).count).to eq 0
+        expect(subject.direct_totals(WEIGHT_SUM_FACET).count).to eq 0
       end
     end
 
@@ -60,9 +61,9 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
         it 'creates no sums for the weight if the issues have 0 weight' do
           subject.assemble_issue_totals
 
-          expect(subject.immediate_totals(COUNT_FACET).count).to eq 1
-          expect(subject.immediate_totals(WEIGHT_SUM_FACET).count).to eq 0
-          expect(subject).to have_immediate_total(ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 1)
+          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)
         end
       end
 
@@ -76,8 +77,8 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
         it 'creates two sums' do
           subject.assemble_issue_totals
 
-          expect(subject).to have_immediate_total(ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 1)
-          expect(subject).to have_immediate_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 2)
+          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)
         end
       end
 
@@ -92,10 +93,10 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
         it 'creates two sums' do
           subject.assemble_issue_totals
 
-          expect(subject).to have_immediate_total(ISSUE_TYPE, COUNT_FACET, OPENED_ISSUE_STATE, 1)
-          expect(subject).to have_immediate_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 2)
-          expect(subject).to have_immediate_total(ISSUE_TYPE, COUNT_FACET, CLOSED_ISSUE_STATE, 3)
-          expect(subject).to have_immediate_total(ISSUE_TYPE, WEIGHT_SUM_FACET, CLOSED_ISSUE_STATE, 5)
+          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
@@ -115,7 +116,7 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
       it 'adds up the number of the child epics' do
         subject.assemble_epic_totals([child_epic_node])
 
-        expect(subject).to have_immediate_total(EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
+        expect(subject).to have_direct_total(EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
       end
     end
   end
@@ -124,8 +125,8 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
     subject { described_class.new(epic_id, [{ parent_id: nil, epic_state_id: CLOSED_EPIC_STATE }]) }
 
     before do
-      allow(subject).to receive(:immediate_totals).with(COUNT_FACET).and_return(immediate_count_totals)
-      allow(subject).to receive(:immediate_totals).with(WEIGHT_SUM_FACET).and_return(immediate_weight_sum_totals)
+      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|
@@ -191,8 +192,8 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
 
       before do
         subject.child_ids << child_epic_id
-        allow(child_epic_node).to receive(:immediate_totals).with(COUNT_FACET).and_return(child_count_totals)
-        allow(child_epic_node).to receive(:immediate_totals).with(WEIGHT_SUM_FACET).and_return(child_weight_sum_totals)
+        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
@@ -210,13 +211,13 @@ describe Gitlab::Graphql::Aggregations::Epics::EpicNode do
         it 'returns the correct count total' do
           subject.calculate_recursive_sums(COUNT_FACET, tree)
 
-          expect(subject.calculated_count_totals.count).to eq 2
+          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.calculated_weight_sum_totals.count).to eq 1
+          expect(subject).to have_calculated_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 2)
         end
       end
     end

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

@@ -108,12 +108,12 @@ describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
           lazy_state = subject.instance_variable_get(:@lazy_state)
           tree = lazy_state[:tree]
 
-          expect(tree[epic_id]).to have_immediate_total(EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
-          expect(tree[epic_id]).to have_immediate_total(ISSUE_TYPE, WEIGHT_SUM_FACET, OPENED_ISSUE_STATE, 5)
-          expect(tree[epic_id]).to have_immediate_total(EPIC_TYPE, COUNT_FACET, CLOSED_EPIC_STATE, 1)
+          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_immediate_total(ISSUE_TYPE, COUNT_FACET, CLOSED_ISSUE_STATE, 4)
-          expect(tree[child_epic_id]).to have_immediate_total(ISSUE_TYPE, WEIGHT_SUM_FACET, CLOSED_ISSUE_STATE, 17)
+          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
@@ -137,8 +137,8 @@ describe Gitlab::Graphql::Aggregations::Epics::LazyEpicAggregate do
           lazy_state = subject.instance_variable_get(:@lazy_state)
           tree = lazy_state[:tree]
 
-          expect(tree[other_epic_id].immediate_count_totals).to be_empty
-          expect(tree[other_epic_id].immediate_weight_sum_totals).to be_empty
+          expect(tree[other_epic_id].direct_count_totals).to be_empty
+          expect(tree[other_epic_id].direct_weight_sum_totals).to be_empty
         end
       end
     end

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

@@ -3,11 +3,7 @@
 require 'spec_helper'
 
 describe Gitlab::Graphql::Loaders::BulkEpicAggregateLoader do
-  let_it_be(:closed_issue_state_id) { Issue.available_states[:closed] }
-  let_it_be(:opened_issue_state_id) { Issue.available_states[:opened] }
-
-  let_it_be(:closed_epic_state_id) { Epic.available_states[:closed] }
-  let_it_be(:opened_epic_state_id) { Epic.available_states[:opened] }
+  include_context 'includes EpicAggregate constants'
 
   let_it_be(:group) { create(:group, :public) }
   let_it_be(:subgroup) { create(:group, :private, parent: group)}
@@ -65,12 +61,12 @@ describe Gitlab::Graphql::Loaders::BulkEpicAggregateLoader do
     it 'sums all the weights, even confidential, or in private groups' do
       expected_result = {
         parent_epic.id => [
-            result_for(parent_epic, issues_state: opened_issue_state_id, issues_count: 5, issues_weight_sum: 4),
-            result_for(parent_epic, issues_state: closed_issue_state_id, issues_count: 1, issues_weight_sum: 1)
+            result_for(parent_epic, issues_state: OPENED_ISSUE_STATE, issues_count: 5, issues_weight_sum: 4),
+            result_for(parent_epic, issues_state: CLOSED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 1)
           ],
         epic_with_issues.id => [
-          result_for(epic_with_issues, issues_state: opened_issue_state_id, issues_count: 5, issues_weight_sum: 4),
-          result_for(epic_with_issues, issues_state: closed_issue_state_id, issues_count: 1, issues_weight_sum: 1)
+          result_for(epic_with_issues, issues_state: OPENED_ISSUE_STATE, issues_count: 5, issues_weight_sum: 4),
+          result_for(epic_with_issues, issues_state: CLOSED_ISSUE_STATE, issues_count: 1, issues_weight_sum: 1)
         ],
         epic_without_issues.id => [
           result_for(epic_without_issues, issues_state: nil, issues_count: 0, issues_weight_sum: 0)
@@ -132,6 +128,6 @@ describe Gitlab::Graphql::Loaders::BulkEpicAggregateLoader do
   end
 
   def result_for(epic, issues_state:, issues_count:, issues_weight_sum:)
-    { 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
 
-RSpec::Matchers.define :have_immediate_total do |type, facet, state, expected_value|
-  match do |epic_node_result|
-    expect(epic_node_result).not_to be_nil
-    immediate_totals = epic_node_result.immediate_totals(facet)
-    expect(immediate_totals).not_to be_empty
-
-    matching = immediate_totals.select { |sum| sum.type == type && sum.facet == facet && sum.state == state && sum.value == expected_value }
-    expect(matching).not_to be_empty
-  end
+%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
-      immediate_totals = epic_node_result.immediate_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 #{immediate_totals.count} immediate sum objects#{", none of which match" if immediate_totals.count > 0}.
-        Sums: #{immediate_totals.inspect}
-      FAILURE_MSG
+    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