Examples with GetDAGStatusRequestProto org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolRPC.GetDAGStatusRequestProto used on opensource projects

Example 1 with GetDAGStatusRequestProto

use of org.apache.tez.dag.api.client.rpc.DAGClientAMProtocolRPC.GetDAGStatusRequestProto in project tez by apache.

the class TestDAGClient method createMockProxy.

private DAGClientAMProtocolBlockingPB createMockProxy(final DAGStatusStateProto stateProto, final long timeout) throws ServiceException {
    DAGClientAMProtocolBlockingPB mock = mock(DAGClientAMProtocolBlockingPB.class);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            GetDAGStatusRequestProto request = (GetDAGStatusRequestProto) invocation.getArguments()[1];
            long sleepTime = request.getTimeout();
            if (timeout != -1) {
                sleepTime = timeout;
            }
            Thread.sleep(sleepTime);
            return GetDAGStatusResponseProto.newBuilder().setDagStatus(constructDagStatusProto(stateProto)).build();
        }
    }).when(mock).getDAGStatus(isNull(RpcController.class), any(GetDAGStatusRequestProto.class));
    return mock;
}