use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.
the class SdkHarnessClientTest method handleCleanupWithStateWhenProcessingBundleFails.
@Test
public void handleCleanupWithStateWhenProcessingBundleFails() throws Exception {
Exception testException = new Exception();
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
StateDelegator mockStateDelegator = mock(StateDelegator.class);
StateDelegator.Registration mockStateRegistration = mock(StateDelegator.Registration.class);
when(mockStateDelegator.registerForProcessBundleInstructionId(any(), any())).thenReturn(mockStateRegistration);
StateRequestHandler mockStateHandler = mock(StateRequestHandler.class);
when(mockStateHandler.getCacheTokens()).thenReturn(Collections.emptyList());
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)), mockStateDelegator);
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
try {
try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockStateHandler, mockProgressHandler)) {
processBundleResponseFuture.completeExceptionally(testException);
}
fail("Exception expected");
} catch (ExecutionException e) {
assertEquals(testException, e.getCause());
verify(mockStateRegistration).abort();
verify(mockOutputReceiver).cancel();
verifyNoMoreInteractions(mockStateRegistration, mockOutputReceiver);
}
}
use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.
the class SdkHarnessClientTest method handleCleanupWhenInputSenderFails.
@Test
public void handleCleanupWhenInputSenderFails() throws Exception {
Exception testException = new Exception();
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))).thenReturn(processBundleResponseFuture);
FullWindowedValueCoder<String> coder = FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
BundleProcessor processor = sdkHarnessClient.getProcessor(descriptor, Collections.singletonList(RemoteInputDestination.of((FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM)));
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
doThrow(testException).when(mockInputSender).close();
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
try {
try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockProgressHandler)) {
// We shouldn't be required to complete the process bundle response future.
}
fail("Exception expected");
} catch (Exception e) {
assertEquals(testException, e);
verify(mockOutputReceiver).cancel();
verifyNoMoreInteractions(mockOutputReceiver);
}
}
use of org.apache.beam.sdk.fn.data.InboundDataClient in project beam by apache.
the class GrpcDataServiceTest method testMultipleClientsSendMessagesAreDirectedToProperConsumers.
@Test
public void testMultipleClientsSendMessagesAreDirectedToProperConsumers() throws Exception {
final LinkedBlockingQueue<BeamFnApi.Elements> clientInboundElements = new LinkedBlockingQueue<>();
ExecutorService executorService = Executors.newCachedThreadPool();
final CountDownLatch waitForInboundElements = new CountDownLatch(1);
GrpcDataService service = GrpcDataService.create(PipelineOptionsFactory.create(), Executors.newCachedThreadPool(), OutboundObserverFactory.serverDirect());
try (GrpcFnServer<GrpcDataService> server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
Collection<Future<Void>> clientFutures = new ArrayList<>();
for (int i = 0; i < 3; ++i) {
final String instructionId = Integer.toString(i);
clientFutures.add(executorService.submit(() -> {
ManagedChannel channel = InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).build();
StreamObserver<Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
outboundObserver.onNext(elementsWithData(instructionId));
waitForInboundElements.await();
outboundObserver.onCompleted();
return null;
}));
}
List<Collection<WindowedValue<String>>> serverInboundValues = new ArrayList<>();
Collection<InboundDataClient> readFutures = new ArrayList<>();
for (int i = 0; i < 3; ++i) {
final Collection<WindowedValue<String>> serverInboundValue = new ArrayList<>();
serverInboundValues.add(serverInboundValue);
readFutures.add(service.receive(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER, serverInboundValue::add));
}
for (InboundDataClient readFuture : readFutures) {
readFuture.awaitCompletion();
}
waitForInboundElements.countDown();
for (Future<Void> clientFuture : clientFutures) {
clientFuture.get();
}
for (int i = 0; i < 3; ++i) {
assertThat(serverInboundValues.get(i), contains(WindowedValue.valueInGlobalWindow("A" + i), WindowedValue.valueInGlobalWindow("B" + i), WindowedValue.valueInGlobalWindow("C" + i)));
}
assertThat(clientInboundElements, empty());
}
}
Aggregations