public GetWorkBudget totalCurrentActiveGetWorkBudget() {
    return computationCache.asMap().values().stream()
        .map(ComputationState::getActiveWorkBudget)
        .reduce(GetWorkBudget.noBudget(), GetWorkBudget::apply);
  }
@Test
  public void testTotalCurrentActiveGetWorkBudget() {
    String computationId = "computationId";
    String computationId2 = "computationId2";
    MapTask mapTask = new MapTask().setStageName("stageName").setSystemName("systemName");
    Map<String, String> userTransformToStateFamilyName =
        ImmutableMap.of("userTransformName", "stateFamilyName");
    ComputationConfig computationConfig =
        ComputationConfig.create(mapTask, userTransformToStateFamilyName, ImmutableMap.of());
    when(configFetcher.fetchConfig(eq(computationId))).thenReturn(Optional.of(computationConfig));
    when(configFetcher.fetchConfig(eq(computationId2))).thenReturn(Optional.of(computationConfig));
    ShardedKey shardedKey = ShardedKey.create(ByteString.EMPTY, 1);
    ShardedKey shardedKey2 = ShardedKey.create(ByteString.EMPTY, 2);

    ExecutableWork work1 = createWork(shardedKey, 1, 1);
    ExecutableWork work2 = createWork(shardedKey2, 2, 2);
    ExecutableWork work3 = createWork(shardedKey2, 3, 3);

    // Activate 1 Work for computationId
    Optional<ComputationState> maybeComputationState = computationStateCache.get(computationId);
    assertTrue(maybeComputationState.isPresent());
    ComputationState computationState = maybeComputationState.get();
    computationState.activateWork(work1);

    // Activate 2 Work(s) for computationId2
    Optional<ComputationState> maybeComputationState2 = computationStateCache.get(computationId);
    assertTrue(maybeComputationState2.isPresent());
    ComputationState computationState2 = maybeComputationState2.get();
    computationState2.activateWork(work2);
    computationState2.activateWork(work3);

    // GetWorkBudget should have 3 items. 1 from computationId, 2 from computationId2.
    assertThat(computationStateCache.totalCurrentActiveGetWorkBudget())
        .isEqualTo(
            GetWorkBudget.builder()
                .setItems(3)
                .setBytes(
                    work1.getWorkItem().getSerializedSize()
                        + work2.getWorkItem().getSerializedSize()
                        + work3.getWorkItem().getSerializedSize())
                .build());
  }