doAnswer(new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
BlockTransferMessage message = BlockTransferMessage.Decoder.fromByteArray(
(byte[]) invocationOnMock.getArguments()[0]);
RpcResponseCallback callback = (RpcResponseCallback) invocationOnMock.getArguments()[1];
callback.onSuccess(new StreamHandle(123, blocks.size()).toByteArray());
assertEquals(new OpenBlocks("app-id", "exec-id", blockIds), message);
return null;
}
}).when(client).sendRpc((byte[]) any(), (RpcResponseCallback) any());
// Respond to each chunk request with a single buffer from our blocks array.
final AtomicInteger expectedChunkIndex = new AtomicInteger(0);
final Iterator<ManagedBuffer> blockIterator = blocks.values().iterator();
doAnswer(new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
try {
long streamId = (Long) invocation.getArguments()[0];
int myChunkIndex = (Integer) invocation.getArguments()[1];
assertEquals(123, streamId);
assertEquals(expectedChunkIndex.getAndIncrement(), myChunkIndex);
ChunkReceivedCallback callback = (ChunkReceivedCallback) invocation.getArguments()[2];
ManagedBuffer result = blockIterator.next();
if (result != null) {
callback.onSuccess(myChunkIndex, result);
} else {
callback.onFailure(myChunkIndex, new RuntimeException("Failed " + myChunkIndex));
}
} catch (Exception e) {
e.printStackTrace();
fail("Unexpected failure");
}