BeamFnControlClient { public void processInstructionRequests(Executor executor) throws InterruptedException, ExecutionException { BeamFnApi.InstructionRequest request; while (!Objects.equals((request = bufferedInstructions.take()), POISON_PILL)) { BeamFnApi.InstructionRequest currentRequest = request; executor.execute( () -> { try { BeamFnApi.InstructionResponse response = delegateOnInstructionRequestType(currentRequest); sendInstructionResponse(response); } catch (Error e) { sendErrorResponse(e); throw e; } }); } onFinish.get(); } BeamFnControlClient(
      String id,
      ApiServiceDescriptor apiServiceDescriptor,
      ManagedChannelFactory channelFactory,
      OutboundObserverFactory outboundObserverFactory,
      EnumMap<
              BeamFnApi.InstructionRequest.RequestCase,
              ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>>
          handlers);  BeamFnControlClient(
      String id,
      BeamFnControlGrpc.BeamFnControlStub controlStub,
      OutboundObserverFactory outboundObserverFactory,
      EnumMap<
              BeamFnApi.InstructionRequest.RequestCase,
              ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>>
          handlers); void processInstructionRequests(Executor executor); BeamFnApi.InstructionResponse delegateOnInstructionRequestType(
      BeamFnApi.InstructionRequest value); void sendInstructionResponse(BeamFnApi.InstructionResponse value);  }
@Test public void testDelegation() throws Exception { AtomicBoolean clientClosedStream = new AtomicBoolean(); BlockingQueue<BeamFnApi.InstructionResponse> values = new LinkedBlockingQueue<>(); BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>(); CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.withOnNext(values::add) .withOnCompleted(() -> clientClosedStream.set(true)) .build(); Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()) .build(); Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()) .addService( new BeamFnControlGrpc.BeamFnControlImplBase() { @Override public StreamObserver<BeamFnApi.InstructionResponse> control( StreamObserver<BeamFnApi.InstructionRequest> outboundObserver) { Uninterruptibles.putUninterruptibly(outboundServerObservers, outboundObserver); return inboundServerObserver; } }) .build(); server.start(); try { EnumMap< BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class); handlers.put( BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, value -> BeamFnApi.InstructionResponse.newBuilder() .setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance())); handlers.put( BeamFnApi.InstructionRequest.RequestCase.REGISTER, value -> { throw FAILURE; }); BeamFnControlClient client = new BeamFnControlClient( "", apiServiceDescriptor, InProcessManagedChannelFactory.create(), OutboundObserverFactory.trivial(), handlers); StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take(); ExecutorService executor = Executors.newCachedThreadPool(); Future<Void> future = executor.submit( () -> { client.processInstructionRequests(executor); return null; }); outboundServerObserver.onNext(SUCCESSFUL_REQUEST); assertEquals(SUCCESSFUL_RESPONSE, values.take()); outboundServerObserver.onNext(UNKNOWN_HANDLER_REQUEST); assertEquals(UNKNOWN_HANDLER_RESPONSE, values.take()); outboundServerObserver.onNext(FAILURE_REQUEST); assertEquals(FAILURE_RESPONSE, values.take()); outboundServerObserver.onCompleted(); future.get(); } finally { server.shutdownNow(); } }
@Test public void testJavaErrorResponse() throws Exception { BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>(); BlockingQueue<Throwable> error = new LinkedBlockingQueue<>(); CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.<BeamFnApi.InstructionResponse>withOnNext( response -> fail(String.format("Unexpected Response %s", response))) .withOnError(error::add) .build(); Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()) .build(); Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()) .addService( new BeamFnControlGrpc.BeamFnControlImplBase() { @Override public StreamObserver<BeamFnApi.InstructionResponse> control( StreamObserver<BeamFnApi.InstructionRequest> outboundObserver) { Uninterruptibles.putUninterruptibly(outboundServerObservers, outboundObserver); return inboundServerObserver; } }) .build(); server.start(); try { EnumMap< BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class); handlers.put( BeamFnApi.InstructionRequest.RequestCase.REGISTER, value -> { throw new Error("Test Error"); }); BeamFnControlClient client = new BeamFnControlClient( "", apiServiceDescriptor, InProcessManagedChannelFactory.create(), OutboundObserverFactory.trivial(), handlers); StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take(); ExecutorService executor = Executors.newCachedThreadPool(); Future<Void> future = executor.submit( () -> { client.processInstructionRequests(executor); return null; }); outboundServerObserver.onNext( InstructionRequest.newBuilder() .setInstructionId("0") .setRegister(RegisterRequest.getDefaultInstance()) .build()); assertThat(error.take(), not(nullValue())); try { future.get(); throw new IllegalStateException("The future should have terminated with an error"); } catch (ExecutionException errorWrapper) { assertThat(errorWrapper.getCause().getMessage(), containsString("Test Error")); } } finally { server.shutdownNow(); } }