BeamFnDataGrpcClient implements BeamFnDataClient { @Override public <T> InboundDataClient receive( ApiServiceDescriptor apiServiceDescriptor, LogicalEndpoint inputLocation, Coder<T> coder, FnDataReceiver<T> consumer) { LOG.debug("Registering consumer for {}", inputLocation); BeamFnDataGrpcMultiplexer client = getClientFor(apiServiceDescriptor); BeamFnDataInboundObserver<T> inboundObserver = BeamFnDataInboundObserver.forConsumer(inputLocation, coder, consumer); client.registerConsumer(inputLocation, inboundObserver); return inboundObserver; } BeamFnDataGrpcClient(
      PipelineOptions options,
      Function<Endpoints.ApiServiceDescriptor, ManagedChannel> channelFactory,
      OutboundObserverFactory outboundObserverFactory); @Override InboundDataClient receive(
      ApiServiceDescriptor apiServiceDescriptor,
      LogicalEndpoint inputLocation,
      Coder<T> coder,
      FnDataReceiver<T> consumer); @Override CloseableFnDataReceiver<T> send(
      Endpoints.ApiServiceDescriptor apiServiceDescriptor,
      LogicalEndpoint outputLocation,
      Coder<T> coder);  }
@Test public void testForInboundConsumer() throws Exception { CountDownLatch waitForClientToConnect = new CountDownLatch(1); Collection<WindowedValue<String>> inboundValuesA = new ConcurrentLinkedQueue<>(); Collection<WindowedValue<String>> inboundValuesB = new ConcurrentLinkedQueue<>(); Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>(); AtomicReference<StreamObserver<BeamFnApi.Elements>> outboundServerObserver = new AtomicReference<>(); CallStreamObserver<BeamFnApi.Elements> inboundServerObserver = TestStreams.withOnNext(inboundServerValues::add).build(); Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()) .build(); Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()) .addService( new BeamFnDataGrpc.BeamFnDataImplBase() { @Override public StreamObserver<BeamFnApi.Elements> data( StreamObserver<BeamFnApi.Elements> outboundObserver) { outboundServerObserver.set(outboundObserver); waitForClientToConnect.countDown(); return inboundServerObserver; } }) .build(); server.start(); try { ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient( PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial()); InboundDataClient readFutureA = clientFactory.receive(apiServiceDescriptor, ENDPOINT_A, CODER, inboundValuesA::add); waitForClientToConnect.await(); outboundServerObserver.get().onNext(ELEMENTS_A_1); outboundServerObserver.get().onNext(ELEMENTS_B_1); Thread.sleep(100); InboundDataClient readFutureB = clientFactory.receive(apiServiceDescriptor, ENDPOINT_B, CODER, inboundValuesB::add); readFutureB.awaitCompletion(); assertThat(inboundValuesB, contains(valueInGlobalWindow("JKL"), valueInGlobalWindow("MNO"))); outboundServerObserver.get().onNext(ELEMENTS_A_2); readFutureA.awaitCompletion(); assertThat( inboundValuesA, contains( valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF"), valueInGlobalWindow("GHI"))); } finally { server.shutdownNow(); } }
@Test public void testForInboundConsumerThatThrows() throws Exception { CountDownLatch waitForClientToConnect = new CountDownLatch(1); AtomicInteger consumerInvoked = new AtomicInteger(); Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>(); AtomicReference<StreamObserver<BeamFnApi.Elements>> outboundServerObserver = new AtomicReference<>(); CallStreamObserver<BeamFnApi.Elements> inboundServerObserver = TestStreams.withOnNext(inboundServerValues::add).build(); Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder() .setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()) .build(); Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()) .addService( new BeamFnDataGrpc.BeamFnDataImplBase() { @Override public StreamObserver<BeamFnApi.Elements> data( StreamObserver<BeamFnApi.Elements> outboundObserver) { outboundServerObserver.set(outboundObserver); waitForClientToConnect.countDown(); return inboundServerObserver; } }) .build(); server.start(); RuntimeException exceptionToThrow = new RuntimeException("TestFailure"); try { ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build(); BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient( PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial()); InboundDataClient readFuture = clientFactory.receive( apiServiceDescriptor, ENDPOINT_A, CODER, t -> { consumerInvoked.incrementAndGet(); throw exceptionToThrow; }); waitForClientToConnect.await(); outboundServerObserver.get().onNext(ELEMENTS_A_1); outboundServerObserver.get().onNext(ELEMENTS_A_2); try { readFuture.awaitCompletion(); fail("Expected channel to fail"); } catch (ExecutionException e) { assertEquals(exceptionToThrow, e.getCause()); } assertThat(inboundServerValues, empty()); assertEquals(1, consumerInvoked.get()); } finally { server.shutdownNow(); } }