Search in sources :

Example 6 with CloseableFnDataReceiver

use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver in project beam by apache.

the class SdkHarnessClientTest method handleCleanupWithStateWhenAwaitingOnClosingOutputReceivers.

@Test
public void handleCleanupWithStateWhenAwaitingOnClosingOutputReceivers() 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(mockOutputReceiver).awaitCompletion();
    RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
    try {
        try (RemoteBundle activeBundle = processor.newBundle(ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mockRemoteOutputReceiver), mockStateHandler, 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);
    }
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) StateRequestHandler(org.apache.beam.runners.fnexecution.state.StateRequestHandler) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) ExpectedException(org.junit.rules.ExpectedException) ExecutionException(java.util.concurrent.ExecutionException) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) CompletableFuture(java.util.concurrent.CompletableFuture) ProcessBundleResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) StateDelegator(org.apache.beam.runners.fnexecution.state.StateDelegator) Test(org.junit.Test)

Example 7 with CloseableFnDataReceiver

use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver 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);
    }
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) StateRequestHandler(org.apache.beam.runners.fnexecution.state.StateRequestHandler) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) ExpectedException(org.junit.rules.ExpectedException) ExecutionException(java.util.concurrent.ExecutionException) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) CompletableFuture(java.util.concurrent.CompletableFuture) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) StateDelegator(org.apache.beam.runners.fnexecution.state.StateDelegator) ExecutionException(java.util.concurrent.ExecutionException) Test(org.junit.Test)

Example 8 with CloseableFnDataReceiver

use of org.apache.beam.sdk.fn.data.CloseableFnDataReceiver 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);
    }
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) ExpectedException(org.junit.rules.ExpectedException) ExecutionException(java.util.concurrent.ExecutionException) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) CompletableFuture(java.util.concurrent.CompletableFuture) BundleProcessor(org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor) Test(org.junit.Test)

Aggregations

CompletableFuture (java.util.concurrent.CompletableFuture)8 InstructionResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse)8 BundleProcessor (org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor)8 CloseableFnDataReceiver (org.apache.beam.sdk.fn.data.CloseableFnDataReceiver)8 InboundDataClient (org.apache.beam.sdk.fn.data.InboundDataClient)8 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)8 Test (org.junit.Test)8 ExecutionException (java.util.concurrent.ExecutionException)6 ExpectedException (org.junit.rules.ExpectedException)6 ProcessBundleResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse)4 StateDelegator (org.apache.beam.runners.fnexecution.state.StateDelegator)3 StateRequestHandler (org.apache.beam.runners.fnexecution.state.StateRequestHandler)3 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)2 ActiveBundle (org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor.ActiveBundle)2