BeamFnStateGrpcClientCache { public BeamFnStateClient forApiServiceDescriptor(ApiServiceDescriptor apiServiceDescriptor) throws IOException { return cache.computeIfAbsent(apiServiceDescriptor, this::createBeamFnStateClient); } BeamFnStateGrpcClientCache(
      IdGenerator idGenerator,
      Function<Endpoints.ApiServiceDescriptor, ManagedChannel> channelFactory,
      OutboundObserverFactory outboundObserverFactory); BeamFnStateClient forApiServiceDescriptor(ApiServiceDescriptor apiServiceDescriptor);  }
@Test public void testCachingOfClient() throws Exception { assertSame( clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(apiServiceDescriptor)); assertNotSame( clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(Endpoints.ApiServiceDescriptor.getDefaultInstance())); }
@Test public void testRequestResponses() throws Exception { BeamFnStateClient client = clientCache.forApiServiceDescriptor(apiServiceDescriptor); CompletableFuture<StateResponse> successfulResponse = new CompletableFuture<>(); CompletableFuture<StateResponse> unsuccessfulResponse = new CompletableFuture<>(); client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), successfulResponse); client.handle(StateRequest.newBuilder().setInstructionId(FAIL), unsuccessfulResponse); StreamObserver<StateResponse> outboundServerObserver = outboundServerObservers.take(); outboundServerObserver.onNext(StateResponse.newBuilder().setId("UNKNOWN ID").build()); handleServerRequest(outboundServerObserver, values.take()); handleServerRequest(outboundServerObserver, values.take()); assertNotNull(successfulResponse.get()); try { unsuccessfulResponse.get(); fail("Expected unsuccessful response"); } catch (ExecutionException e) { assertThat(e.toString(), containsString(TEST_ERROR)); } }
@Test public void testServerErrorCausesPendingAndFutureCallsToFail() throws Exception { BeamFnStateClient client = clientCache.forApiServiceDescriptor(apiServiceDescriptor); CompletableFuture<StateResponse> inflight = new CompletableFuture<>(); client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), inflight); StreamObserver<StateResponse> outboundServerObserver = outboundServerObservers.take(); outboundServerObserver.onError( new StatusRuntimeException(Status.INTERNAL.withDescription(SERVER_ERROR))); try { inflight.get(); fail("Expected unsuccessful response due to server error"); } catch (ExecutionException e) { assertThat(e.toString(), containsString(SERVER_ERROR)); } CompletableFuture<StateResponse> late = new CompletableFuture<>(); client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), late); try { inflight.get(); fail("Expected unsuccessful response due to server error"); } catch (ExecutionException e) { assertThat(e.toString(), containsString(SERVER_ERROR)); } }
@Test public void testServerCompletionCausesPendingAndFutureCallsToFail() throws Exception { BeamFnStateClient client = clientCache.forApiServiceDescriptor(apiServiceDescriptor); CompletableFuture<StateResponse> inflight = new CompletableFuture<>(); client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), inflight); StreamObserver<StateResponse> outboundServerObserver = outboundServerObservers.take(); outboundServerObserver.onCompleted(); try { inflight.get(); fail("Expected unsuccessful response due to server completion"); } catch (ExecutionException e) { assertThat(e.toString(), containsString("Server hanged up")); } CompletableFuture<StateResponse> late = new CompletableFuture<>(); client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS), late); try { inflight.get(); fail("Expected unsuccessful response due to server completion"); } catch (ExecutionException e) { assertThat(e.toString(), containsString("Server hanged up")); } }