throws IOException {
TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
BlockFetchStarter fetchStarter = mock(BlockFetchStarter.class);
Stubber stub = null;
// Contains all blockIds that are referenced across all interactions.
final LinkedHashSet<String> blockIds = Sets.newLinkedHashSet();
for (final Map<String, Object> interaction : interactions) {
blockIds.addAll(interaction.keySet());
Answer<Void> answer = new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
try {
// Verify that the RetryingBlockFetcher requested the expected blocks.
String[] requestedBlockIds = (String[]) invocationOnMock.getArguments()[0];
String[] desiredBlockIds = interaction.keySet().toArray(new String[interaction.size()]);
assertArrayEquals(desiredBlockIds, requestedBlockIds);
// Now actually invoke the success/failure callbacks on each block.
BlockFetchingListener retryListener =
(BlockFetchingListener) invocationOnMock.getArguments()[1];
for (Map.Entry<String, Object> block : interaction.entrySet()) {
String blockId = block.getKey();
Object blockValue = block.getValue();
if (blockValue instanceof ManagedBuffer) {
retryListener.onBlockFetchSuccess(blockId, (ManagedBuffer) blockValue);
} else if (blockValue instanceof Exception) {
retryListener.onBlockFetchFailure(blockId, (Exception) blockValue);
} else {
fail("Can only handle ManagedBuffers and Exceptions, got " + blockValue);
}
}
return null;
} catch (Throwable e) {
e.printStackTrace();
throw e;
}
}
};
// This is either the first stub, or should be chained behind the prior ones.
if (stub == null) {
stub = doAnswer(answer);
} else {
stub.doAnswer(answer);
}
}
assert stub != null;
stub.when(fetchStarter).createAndStart((String[]) any(), (BlockFetchingListener) anyObject());
String[] blockIdArray = blockIds.toArray(new String[blockIds.size()]);
new RetryingBlockFetcher(conf, fetchStarter, blockIdArray, listener).start();
}