use of io.grpc.ForwardingClientCall.SimpleForwardingClientCall in project beam by apache.
the class GrpcContextHeaderAccessorProviderTest method testWorkerIdOnConnect.
@SuppressWarnings("unchecked")
@Test
public void testWorkerIdOnConnect() throws Exception {
final String worker1 = "worker1";
CompletableFuture<String> workerId = new CompletableFuture<>();
Consumer<StreamObserver<Elements>> consumer = elementsStreamObserver -> {
workerId.complete(GrpcContextHeaderAccessorProvider.getHeaderAccessor().getSdkWorkerId());
elementsStreamObserver.onCompleted();
};
TestDataService testService = new TestDataService(Mockito.mock(StreamObserver.class), consumer);
ApiServiceDescriptor serviceDescriptor = ApiServiceDescriptor.newBuilder().setUrl("testServer").build();
cleanupRule.register(InProcessServerFactory.create().create(ImmutableList.of(testService), serviceDescriptor));
final Metadata.Key<String> workerIdKey = Metadata.Key.of("worker_id", Metadata.ASCII_STRING_MARSHALLER);
Channel channel = cleanupRule.register(InProcessChannelBuilder.forName(serviceDescriptor.getUrl()).intercept(new ClientInterceptor() {
@Override
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
ClientCall<ReqT, RespT> call = next.newCall(method, callOptions);
return new SimpleForwardingClientCall<ReqT, RespT>(call) {
@Override
public void start(ClientCall.Listener<RespT> responseListener, Metadata headers) {
headers.put(workerIdKey, worker1);
super.start(responseListener, headers);
}
};
}
}).build());
BeamFnDataGrpc.BeamFnDataStub stub = BeamFnDataGrpc.newStub(channel);
stub.data(Mockito.mock(StreamObserver.class)).onCompleted();
Assert.assertEquals(worker1, workerId.get());
}
Aggregations