Code review comment: We should be checking that the mock FnApiControlClient got a ProcessBundle InstructionRequest with the cache token specified and not verifying that the state request handler was invoked.

Code before review:
public void verifyCacheTokensAreUsedInNewBundleRequest() {
CompletableFuture<InstructionResponse> registerResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(registerResponseFuture);
ProcessBundleDescriptor descriptor1 =
ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputs =
Collections.singletonMap(
"inputPC",
RemoteInputDestination.of(
(FullWindowedValueCoder)
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());
Mockito.verify(stateRequestHandler).getCacheTokens();
}

Line under review: Mockito.verify(stateRequestHandler).getCacheTokens();

Code after addressing the review:
public void verifyCacheTokensAreUsedInNewBundleRequest() {
CompletableFuture<InstructionResponse> registerResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(registerResponseFuture);
ProcessBundleDescriptor descriptor1 =
ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
Map<String, RemoteInputDestination> remoteInputs =
Collections.singletonMap(
"inputPC",
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(2)).handle(reqCaptor.capture());
List<BeamFnApi.InstructionRequest> requests = reqCaptor.getAllValues();
assertThat(
requests.get(0).getRequestCase(), is(BeamFnApi.InstructionRequest.RequestCase.REGISTER));
assertThat(
requests.get(1).getRequestCase(),
is(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE));
assertThat(requests.get(1).getProcessBundle().getCacheTokensList(), is(cacheTokens));
}