Examples with WorkerStatusRequest org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest used on opensource projects

Example 1 with WorkerStatusRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest in project beam by apache.

the class BeamWorkerStatusGrpcServiceTest method testGetWorkerStatusReturnError.

@Test
public void testGetWorkerStatusReturnError() throws Exception {
    StreamObserver<WorkerStatusResponse> observer = stub.workerStatus(mockObserver);
    waitAndGetStatusClient(ID);
    doAnswer((invocation) -> {
        WorkerStatusRequest request = (WorkerStatusRequest) invocation.getArguments()[0];
        observer.onNext(WorkerStatusResponse.newBuilder().setId(request.getId()).setError("error").build());
        return null;
    }).when(mockObserver).onNext(any());
    CompletableFuture<String> future = service.getWorkerStatus(ID);
    String response = future.get(5, TimeUnit.SECONDS);
    assertEquals("error", response);
}

Example 2 with WorkerStatusRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest in project beam by apache.

the class BeamWorkerStatusGrpcServiceTest method testGetWorkerStatusSuccess.

@Test
public void testGetWorkerStatusSuccess() throws Exception {
    StreamObserver<WorkerStatusResponse> observer = stub.workerStatus(mockObserver);
    waitAndGetStatusClient(ID);
    doAnswer((invocation) -> {
        WorkerStatusRequest request = (WorkerStatusRequest) invocation.getArguments()[0];
        observer.onNext(WorkerStatusResponse.newBuilder().setId(request.getId()).setStatusInfo("status").build());
        return null;
    }).when(mockObserver).onNext(any());
    CompletableFuture<String> future = service.getWorkerStatus(ID);
    String response = future.get(5, TimeUnit.SECONDS);
    assertEquals("status", response);
}

Example 3 with WorkerStatusRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest 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));
    }
}

Example 4 with WorkerStatusRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest in project beam by apache.

the class BeamFnStatusClientTest method testWorkerStatusResponse.

@Test
public void testWorkerStatusResponse() throws Exception {
    BlockingQueue<WorkerStatusResponse> values = new LinkedBlockingQueue<>();
    BlockingQueue<StreamObserver<WorkerStatusRequest>> requestObservers = new LinkedBlockingQueue<>();
    StreamObserver<WorkerStatusResponse> inboundServerObserver = TestStreams.withOnNext(values::add).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnWorkerStatusImplBase() {

        @Override
        public StreamObserver<WorkerStatusResponse> workerStatus(StreamObserver<WorkerStatusRequest> responseObserver) {
            Uninterruptibles.putUninterruptibly(requestObservers, responseObserver);
            return inboundServerObserver;
        }
    }).build();
    server.start();
    try {
        BundleProcessorCache processorCache = mock(BundleProcessorCache.class);
        when(processorCache.getActiveBundleProcessors()).thenReturn(Collections.emptyMap());
        ManagedChannelFactory channelFactory = ManagedChannelFactory.createInProcess();
        new BeamFnStatusClient(apiServiceDescriptor, channelFactory::forDescriptor, processorCache, PipelineOptionsFactory.create(), Caches.noop());
        StreamObserver<WorkerStatusRequest> requestObserver = requestObservers.take();
        requestObserver.onNext(WorkerStatusRequest.newBuilder().setId("id").build());
        WorkerStatusResponse response = values.take();
        assertThat(response.getStatusInfo(), containsString("No active processing bundles."));
        assertThat(response.getId(), is("id"));
    } finally {
        server.shutdownNow();
    }
}