Examples with InstructionResponse org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse used on opensource projects

Example 26 with InstructionResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse in project beam by apache.

the class SdkHarnessClientTest method testProgressAndSplitCallsAreIgnoredWhenBundleIsComplete.

@Test
public void testProgressAndSplitCallsAreIgnoredWhenBundleIsComplete() throws Exception {
    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.send(any(), eq(coder))).thenReturn(mock(CloseableFnDataReceiver.class));
    RemoteBundle activeBundle = processor.newBundle(Collections.emptyMap(), BundleProgressHandler.ignored());
    // Correlating the request and response is owned by the underlying
    // FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
    // the response.
    BeamFnApi.ProcessBundleResponse response = ProcessBundleResponse.getDefaultInstance();
    processBundleResponseFuture.complete(BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build());
    activeBundle.close();
    verify(fnApiControlClient).registerProcessBundleDescriptor(any(ProcessBundleDescriptor.class));
    verify(fnApiControlClient).handle(any(BeamFnApi.InstructionRequest.class));
    activeBundle.requestProgress();
    activeBundle.split(0);
    verifyNoMoreInteractions(fnApiControlClient);
}

Example 27 with InstructionResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse 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);
    }
}