use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver in project beam by apache.
the class SdkHarnessClientTest method handleCleanupWhenAwaitingOnClosingOutputReceivers.
@Test
public void handleCleanupWhenAwaitingOnClosingOutputReceivers() 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(mockOutputReceiver).awaitCompletion();
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
try {
try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockProgressHandler)) {
// Correlating the ProcessBundleRequest and ProcessBundleResponse is owned by the underlying
// FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
// the response.
//
// Currently there are no fields so there's nothing to check. This test is formulated
// to match the pattern it should have if/when the response is meaningful.
BeamFnApi.ProcessBundleResponse response = BeamFnApi.ProcessBundleResponse.getDefaultInstance();
processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build());
}
fail("Exception expected");
} catch (Exception e) {
assertEquals(testException, e);
}
}
use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver in project beam by apache.
the class SdkHarnessClientTest method handleCleanupWhenProcessingBundleFails.
@Test
public void handleCleanupWhenProcessingBundleFails() 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);
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
try {
try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockProgressHandler)) {
processBundleResponseFuture.completeExceptionally(testException);
}
fail("Exception expected");
} catch (ExecutionException e) {
assertEquals(testException, e.getCause());
verify(mockOutputReceiver).cancel();
verifyNoMoreInteractions(mockOutputReceiver);
}
}
use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver in project beam by apache.
the class SdkHarnessClientTest method handleCleanupWithStateWhenInputSenderFails.
@Test
public void handleCleanupWithStateWhenInputSenderFails() 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);
doThrow(testException).when(mockInputSender).close();
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
try {
try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockStateHandler, mockProgressHandler)) {
// We shouldn't be required to complete the process bundle response future.
}
fail("Exception expected");
} catch (Exception e) {
assertEquals(testException, e);
verify(mockStateRegistration).abort();
verify(mockOutputReceiver).cancel();
verifyNoMoreInteractions(mockStateRegistration, mockOutputReceiver);
}
}
use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver in project beam by apache.
the class SdkHarnessClientTest method testBundleCheckpointCallback.
@Test
public void testBundleCheckpointCallback() throws 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);
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
BundleSplitHandler mockSplitHandler = mock(BundleSplitHandler.class);
BundleCheckpointHandler mockCheckpointHandler = mock(BundleCheckpointHandler.class);
BundleFinalizationHandler mockFinalizationHandler = mock(BundleFinalizationHandler.class);
ProcessBundleResponse response = ProcessBundleResponse.newBuilder().addResidualRoots(DelayedBundleApplication.getDefaultInstance()).build();
try (ActiveBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), Collections.emptyMap(), (request) -> {
throw new UnsupportedOperationException();
}, mockProgressHandler, mockSplitHandler, mockCheckpointHandler, mockFinalizationHandler)) {
processBundleResponseFuture.complete(InstructionResponse.newBuilder().setProcessBundle(response).build());
}
verify(mockProgressHandler).onCompleted(response);
verify(mockCheckpointHandler).onCheckpoint(response);
verifyZeroInteractions(mockFinalizationHandler, mockSplitHandler);
}
use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver in project beam by apache.
the class SdkHarnessClientTest method testBundleFinalizationCallback.
@Test
public void testBundleFinalizationCallback() throws 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);
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
BundleSplitHandler mockSplitHandler = mock(BundleSplitHandler.class);
BundleCheckpointHandler mockCheckpointHandler = mock(BundleCheckpointHandler.class);
BundleFinalizationHandler mockFinalizationHandler = mock(BundleFinalizationHandler.class);
ProcessBundleResponse response = ProcessBundleResponse.newBuilder().setRequiresFinalization(true).build();
String bundleId;
try (ActiveBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), Collections.emptyMap(), (request) -> {
throw new UnsupportedOperationException();
}, mockProgressHandler, mockSplitHandler, mockCheckpointHandler, mockFinalizationHandler)) {
bundleId = activeBundle.getId();
processBundleResponseFuture.complete(InstructionResponse.newBuilder().setProcessBundle(response).build());
}
verify(mockProgressHandler).onCompleted(response);
verify(mockFinalizationHandler).requestsFinalization(bundleId);
verifyZeroInteractions(mockCheckpointHandler, mockSplitHandler);
}
Aggregations