SdkHarnessClient implements AutoCloseable { public BundleProcessor getProcessor( BeamFnApi.ProcessBundleDescriptor descriptor, List<RemoteInputDestination> remoteInputDesinations) { checkState( !descriptor.hasStateApiServiceDescriptor(), "The %s cannot support a %s containing a state %s.", BundleProcessor.class.getSimpleName(), BeamFnApi.ProcessBundleDescriptor.class.getSimpleName(), Endpoints.ApiServiceDescriptor.class.getSimpleName()); return getProcessor(descriptor, remoteInputDesinations, NoOpStateDelegator.INSTANCE); } private  SdkHarnessClient(
      InstructionRequestHandler fnApiControlClient,
      FnDataService fnApiDataService,
      IdGenerator idGenerator); static SdkHarnessClient usingFnApiClient(
      InstructionRequestHandler fnApiControlClient, FnDataService fnApiDataService); SdkHarnessClient withIdGenerator(IdGenerator idGenerator); BundleProcessor getProcessor(
      BeamFnApi.ProcessBundleDescriptor descriptor,
      List<RemoteInputDestination> remoteInputDesinations); BundleProcessor getProcessor(
      BeamFnApi.ProcessBundleDescriptor descriptor,
      List<RemoteInputDestination> remoteInputDesinations,
      StateDelegator stateDelegator); BundleProcessor getProcessor(
      BeamFnApi.ProcessBundleDescriptor descriptor,
      List<RemoteInputDestination> remoteInputDestinations,
      StateDelegator stateDelegator,
      Map<String, Map<String, TimerSpec>> timerSpecs); @Override void close();  }
@Test public void testRegister() throws Exception { ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build(); List<RemoteInputDestination> remoteInputs = Collections.singletonList( RemoteInputDestination.of( (FullWindowedValueCoder) FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM)); BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs); verify(fnApiControlClient).registerProcessBundleDescriptor(descriptor1); }
@Test public void testRegisterCachesBundleProcessors() throws Exception { ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build(); ProcessBundleDescriptor descriptor2 = ProcessBundleDescriptor.newBuilder().setId("descriptor2").build(); List<RemoteInputDestination> remoteInputs = Collections.singletonList( RemoteInputDestination.of( (FullWindowedValueCoder) FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM)); BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs); BundleProcessor processor2 = sdkHarnessClient.getProcessor(descriptor2, remoteInputs); assertNotSame(processor1, processor2); assertSame(processor1, sdkHarnessClient.getProcessor(descriptor1, remoteInputs)); }
@Test public void testRegisterWithStateRequiresStateDelegator() throws Exception { ProcessBundleDescriptor descriptor = ProcessBundleDescriptor.newBuilder() .setId("test") .setStateApiServiceDescriptor(ApiServiceDescriptor.newBuilder().setUrl("foo")) .build(); List<RemoteInputDestination> remoteInputs = Collections.singletonList( RemoteInputDestination.of( (FullWindowedValueCoder) FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM)); thrown.expect(IllegalStateException.class); thrown.expectMessage("containing a state"); sdkHarnessClient.getProcessor(descriptor, remoteInputs); }
@Test public void testNewBundleNoDataDoesNotCrash() 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)); try (RemoteBundle activeBundle = processor.newBundle(Collections.emptyMap(), BundleProgressHandler.ignored())) { BeamFnApi.ProcessBundleResponse response = ProcessBundleResponse.getDefaultInstance(); processBundleResponseFuture.complete( BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build()); } }
@Test public void testNewBundleAndProcessElements() throws Exception { SdkHarnessClient client = harness.client(); BundleProcessor processor = client.getProcessor( descriptor, Collections.singletonList( RemoteInputDestination.of( (FullWindowedValueCoder) FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM))); Collection<WindowedValue<String>> outputs = new ArrayList<>(); try (RemoteBundle activeBundle = processor.newBundle( Collections.singletonMap( SDK_GRPC_WRITE_TRANSFORM, RemoteOutputReceiver.of( FullWindowedValueCoder.of( LengthPrefixCoder.of(StringUtf8Coder.of()), Coder.INSTANCE), outputs::add)), BundleProgressHandler.ignored())) { FnDataReceiver<WindowedValue<?>> bundleInputReceiver = Iterables.getOnlyElement(activeBundle.getInputReceivers().values()); bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("foo")); bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("bar")); bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("baz")); } assertThat( outputs, containsInAnyOrder( WindowedValue.valueInGlobalWindow("spam"), WindowedValue.valueInGlobalWindow("ham"), WindowedValue.valueInGlobalWindow("eggs"))); }
@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); } }
@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); } }
@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)) { BeamFnApi.ProcessBundleResponse response = BeamFnApi.ProcessBundleResponse.getDefaultInstance(); processBundleResponseFuture.complete( BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build()); } fail("Exception expected"); } catch (Exception e) { assertEquals(testException, e); } }
@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)) { BeamFnApi.ProcessBundleResponse response = BeamFnApi.ProcessBundleResponse.getDefaultInstance(); processBundleResponseFuture.complete( BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response).build()); } fail("Exception expected"); } catch (Exception e) { assertEquals(testException, e); } }
@Test public void verifyCacheTokensAreUsedInNewBundleRequest() throws InterruptedException { when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))) .thenReturn( CompletableFuture.<InstructionResponse>completedFuture( InstructionResponse.newBuilder().build())); ProcessBundleDescriptor descriptor1 = ProcessBundleDescriptor.newBuilder().setId("descriptor1").build(); List<RemoteInputDestination> remoteInputs = Collections.singletonList( RemoteInputDestination.of( FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE), SDK_GRPC_READ_TRANSFORM)); BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs); when(dataService.send(any(), any())).thenReturn(mock(CloseableFnDataReceiver.class)); StateRequestHandler stateRequestHandler = Mockito.mock(StateRequestHandler.class); List<BeamFnApi.ProcessBundleRequest.CacheToken> cacheTokens = Collections.singletonList( BeamFnApi.ProcessBundleRequest.CacheToken.newBuilder().getDefaultInstanceForType()); when(stateRequestHandler.getCacheTokens()).thenReturn(cacheTokens); processor1.newBundle( ImmutableMap.of(SDK_GRPC_WRITE_TRANSFORM, mock(RemoteOutputReceiver.class)), stateRequestHandler, BundleProgressHandler.ignored()); ArgumentCaptor<BeamFnApi.InstructionRequest> reqCaptor = ArgumentCaptor.forClass(BeamFnApi.InstructionRequest.class); Mockito.verify(fnApiControlClient, Mockito.times(1)).handle(reqCaptor.capture()); List<BeamFnApi.InstructionRequest> requests = reqCaptor.getAllValues(); assertThat( requests.get(0).getRequestCase(), is(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE)); assertThat(requests.get(0).getProcessBundle().getCacheTokensList(), is(cacheTokens)); }
@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); }
@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); }