Search in sources :

Example 1 with BeamFnWorkerStatusImplBase

use of org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc.BeamFnWorkerStatusImplBase 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();
    }
}
Also used : StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) WorkerStatusResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusResponse) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) BundleProcessorCache(org.apache.beam.fn.harness.control.ProcessBundleHandler.BundleProcessorCache) WorkerStatusRequest(org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) ManagedChannelFactory(org.apache.beam.sdk.fn.channel.ManagedChannelFactory) BeamFnWorkerStatusImplBase(org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc.BeamFnWorkerStatusImplBase) Test(org.junit.Test)

Aggregations

LinkedBlockingQueue (java.util.concurrent.LinkedBlockingQueue)1 BundleProcessorCache (org.apache.beam.fn.harness.control.ProcessBundleHandler.BundleProcessorCache)1 WorkerStatusRequest (org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusRequest)1 WorkerStatusResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.WorkerStatusResponse)1 BeamFnWorkerStatusImplBase (org.apache.beam.model.fnexecution.v1.BeamFnWorkerStatusGrpc.BeamFnWorkerStatusImplBase)1 ManagedChannelFactory (org.apache.beam.sdk.fn.channel.ManagedChannelFactory)1 Server (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)1 StreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver)1 Test (org.junit.Test)1