GrpcContextHeaderAccessorProvider { public static HeaderAccessor getHeaderAccessor() { return HEADER_ACCESSOR; }  static ServerInterceptor interceptor(); static HeaderAccessor getHeaderAccessor();  }
@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()); }
@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()); TestDataService testService = new TestDataService(Mockito.mock(StreamObserver.class), consumer); ApiServiceDescriptor serviceDescriptor = ApiServiceDescriptor.newBuilder().setUrl("testServer").build(); Server server = InProcessServerFactory.create().create(ImmutableList.of(testService), serviceDescriptor); final Metadata.Key<String> workerIdKey = Metadata.Key.of("worker_id", Metadata.ASCII_STRING_MARSHALLER); Channel channel = 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)); server.shutdown(); Assert.assertEquals(worker1, workerId.get()); }