ImmutableClassesGiraphConfiguration conf =
Mockito.mock(ImmutableClassesGiraphConfiguration.class);
Mockito.when(conf.getAggregatorWriterClass()).thenReturn(
TextAggregatorWriter.class);
Progressable progressable = Mockito.mock(Progressable.class);
MasterAggregatorHandler handler =
new MasterAggregatorHandler(conf, progressable);
String regularAggName = "regular";
LongWritable regularValue = new LongWritable(5);
handler.registerAggregator(regularAggName, LongSumAggregator.class);
handler.setAggregatedValue(regularAggName, regularValue);
String persistentAggName = "persistent";
DoubleWritable persistentValue = new DoubleWritable(10.5);
handler.registerPersistentAggregator(persistentAggName,
DoubleOverwriteAggregator.class);
handler.setAggregatedValue(persistentAggName, persistentValue);
for (AggregatorWrapper<Writable> aggregator :
getAggregatorMap(handler).values()) {
aggregator.setPreviousAggregatedValue(
aggregator.getCurrentAggregatedValue());
}
ByteArrayOutputStream out = new ByteArrayOutputStream();
handler.write(new DataOutputStream(out));
MasterAggregatorHandler restartedHandler =
new MasterAggregatorHandler(conf, progressable);
restartedHandler.readFields(
new DataInputStream(new ByteArrayInputStream(out.toByteArray())));
assertEquals(2, getAggregatorMap(restartedHandler).size());
AggregatorWrapper<Writable> regularAgg =
getAggregatorMap(restartedHandler).get(regularAggName);
assertTrue(
regularAgg.getAggregatorClass().equals(LongSumAggregator.class));
assertEquals(regularValue, regularAgg.getPreviousAggregatedValue());
assertEquals(regularValue,
restartedHandler.<LongWritable>getAggregatedValue(regularAggName));
assertFalse(regularAgg.isPersistent());
AggregatorWrapper<Writable> persistentAgg =
getAggregatorMap(restartedHandler).get(persistentAggName);
assertTrue(persistentAgg.getAggregatorClass().equals
(DoubleOverwriteAggregator.class));
assertEquals(persistentValue, persistentAgg.getPreviousAggregatedValue());
assertEquals(persistentValue,
restartedHandler.<LongWritable>getAggregatedValue(persistentAggName));
assertTrue(persistentAgg.isPersistent());
}