bb.putInt(shufflePort);
bb.position(0);
AuxiliaryServiceHelper.setServiceDataIntoEnv(ShuffleUtils.SHUFFLE_HANDLER_SERVICE_ID, bb, auxEnv);
TezOutputContext outputContext = new TezOutputContextImpl(conf,
appAttemptNumber, tezUmbilical, dagName, taskVertexName, destinationVertexName,
taskAttemptID, counters, 0, userPayload, runtimeTask,
null, auxEnv, new MemoryDistributor(1, 1, conf) , mock(OutputDescriptor.class));
List<Event> events = null;
events = kvOutput.initialize(outputContext);
assertTrue(events != null && events.size() == 0);
KeyValueWriter kvWriter = kvOutput.getWriter();
List<KVPair> data = KVDataGen.generateTestData(true);
for (KVPair kvp : data) {
kvWriter.write(kvp.getKey(), kvp.getvalue());
}
events = kvOutput.close();
assertTrue(events != null && events.size() == 1);
DataMovementEvent dmEvent = (DataMovementEvent)events.get(0);
assertEquals("Invalid source index", 0, dmEvent.getSourceIndex());
DataMovementEventPayloadProto shufflePayload = DataMovementEventPayloadProto
.parseFrom(dmEvent.getUserPayload());
assertFalse(shufflePayload.hasEmptyPartitions());
assertEquals(outputContext.getUniqueIdentifier(), shufflePayload.getPathComponent());
assertEquals(shufflePort, shufflePayload.getPort());
assertEquals("host", shufflePayload.getHost());
}