Examples with BeamFnWorkerStatusStub org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc.BeamFnWorkerStatusStub used on opensource projects

Example 1 with BeamFnWorkerStatusStub

use of org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc.BeamFnWorkerStatusStub in project beam by apache.

the class BeamWorkerStatusGrpcServiceTest method testGetAllWorkerStatuses.

@Test
public void testGetAllWorkerStatuses() throws Exception {
    Set<String> ids = Sets.newHashSet("id0", "id3", "id11", "id12", "id21");
    for (String id : ids) {
        StreamObserver<WorkerStatusRequest> requestObserverMock = mock(StreamObserver.class);
        BeamFnWorkerStatusStub workerStatusStub = BeamFnWorkerStatusGrpc.newStub(channel).withInterceptors(AddHarnessIdInterceptor.create(id));
        StreamObserver<WorkerStatusResponse> observer = workerStatusStub.workerStatus(requestObserverMock);
        // wait for connection before proceeding to avoid race condition.
        waitAndGetStatusClient(id);
        doAnswer((invocation) -> {
            WorkerStatusRequest request = (WorkerStatusRequest) invocation.getArguments()[0];
            observer.onNext(WorkerStatusResponse.newBuilder().setId(request.getId()).setStatusInfo("status").build());
            return null;
        }).when(requestObserverMock).onNext(any());
    }
    Map<String, String> allWorkerStatuses = service.getAllWorkerStatuses(5, TimeUnit.SECONDS);
    assertEquals(ids, allWorkerStatuses.keySet());
    for (String id : ids) {
        assertEquals("status", allWorkerStatuses.get(id));
    }
}