use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.
the class ServerFactoryTest method defaultServerWithPortSupplier.
@Test
public void defaultServerWithPortSupplier() throws Exception {
Endpoints.ApiServiceDescriptor apiServiceDescriptor = runTestUsing(ServerFactory.createWithPortSupplier(() -> 65535), ManagedChannelFactory.createDefault());
HostAndPort hostAndPort = HostAndPort.fromString(apiServiceDescriptor.getUrl());
assertThat(hostAndPort.getHost(), anyOf(equalTo(InetAddress.getLoopbackAddress().getHostName()), equalTo(InetAddress.getLoopbackAddress().getHostAddress())));
assertThat(hostAndPort.getPort(), is(65535));
}
use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor 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());
}
use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.
the class ProcessBundleDescriptors method addStageInput.
private static RemoteInputDestination<WindowedValue<?>> addStageInput(ApiServiceDescriptor dataEndpoint, PCollectionNode inputPCollection, Components.Builder components, WireCoderSetting wireCoderSetting) throws IOException {
String inputWireCoderId = WireCoders.addSdkWireCoder(inputPCollection, components, wireCoderSetting);
@SuppressWarnings("unchecked") Coder<WindowedValue<?>> wireCoder = (Coder) WireCoders.instantiateRunnerWireCoder(inputPCollection, components.build(), wireCoderSetting);
RemoteGrpcPort inputPort = RemoteGrpcPort.newBuilder().setApiServiceDescriptor(dataEndpoint).setCoderId(inputWireCoderId).build();
String inputId = uniqueId(String.format("fn/read/%s", inputPCollection.getId()), components::containsTransforms);
PTransform inputTransform = RemoteGrpcPortRead.readFromPort(inputPort, inputPCollection.getId()).toPTransform();
components.putTransforms(inputId, inputTransform);
return RemoteInputDestination.of(wireCoder, inputId);
}
Aggregations