@Override
  public Map<String, Integer> getCounts(UUID jobId) {
    Query<Entity> query = getCountsQuery(jobId);
    QueryResults<Entity> results = datastore.run(query);
    ImmutableMap.Builder<String, Integer> countsMapBuilder = ImmutableMap.builder();

    while (results.hasNext()) {
      Entity result = results.next();
      String dataType = result.getKey().getName();
      long count = result.getLong(COUNTS_FIELD);
      countsMapBuilder.put(dataType, (int) count);
    }

    return countsMapBuilder.build();
  }
@Test
  public void canAddExistingKeysToCurrentCountsTest() throws IOException {
    addItemToJobStoreCounts(ITEM_NAME);
    addItemToJobStoreCounts(ITEM_NAME);

    final Map<String, Integer> counts = googleJobStore.getCounts(JOB_ID);
    Truth.assertThat(counts.size()).isEqualTo(1);
    Truth.assertThat(counts.get(ITEM_NAME)).isEqualTo(2);
  }