@Test
public void testAggregation() throws Exception {

  // Test aggregation with multiple groups.
  int groups = 3;
  int n = 50;
  TimelineEntities testEntities = generateTestEntities(groups, n);
  TimelineEntity resultEntity = TimelineCollector.aggregateEntities(
      testEntities, "test_result", "TEST_AGGR", true);
  assertThat(resultEntity.getMetrics()).hasSize(groups * 3);
  for (int i = 0; i < groups; i++) {
    Set<TimelineMetric> metrics = resultEntity.getMetrics();
    for (TimelineMetric m : metrics) {
      if (m.getId().startsWith("HDFS_BYTES_WRITE")) {
        assertEquals(100 * n, m.getSingleDataValue().intValue());
      } else if (m.getId().startsWith("VCORES_USED")) {
        assertEquals(3 * n, m.getSingleDataValue().intValue());
      } else if (m.getId().startsWith("TXN_FINISH_TIME")) {
        assertEquals(n - 1, m.getSingleDataValue());
      } else {
        fail("Unrecognized metric! " + m.getId());
      }
    }
  }
  // Test aggregation with a single group.
  TimelineEntities testEntities1 = generateTestEntities(1, n);
  TimelineEntity resultEntity1 = TimelineCollector.aggregateEntities(
      testEntities1, "test_result", "TEST_AGGR", false);
  assertThat(resultEntity1.getMetrics()).hasSize(3);
  Set<TimelineMetric> metrics = resultEntity1.getMetrics();
  for (TimelineMetric m : metrics) {
    if (m.getId().equals("HDFS_BYTES_WRITE")) {
      assertEquals(100 * n, m.getSingleDataValue().intValue());
    } else if (m.getId().equals("VCORES_USED")) {
      assertEquals(3 * n, m.getSingleDataValue().intValue());
    } else if (m.getId().equals("TXN_FINISH_TIME")) {
      assertEquals(n - 1, m.getSingleDataValue());
    } else {
      fail("Unrecognized metric! " + m.getId());
    }
  }
}