use of org.apache.beam.fn.harness.Cache in project beam by apache.
the class ProcessBundleHandlerTest method testBundleProcessorIsResetWhenAddedBackToCache.
@Test
public void testBundleProcessorIsResetWhenAddedBackToCache() throws Exception {
BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).build();
Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateGrpcClientCache */
null, /* finalizeBundleHandler */
new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, (context) -> null), Caches.noop(), new TestBundleProcessorCache());
assertThat(TestBundleProcessor.resetCnt, equalTo(0));
handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("998L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
// Check that BundleProcessor is reset when added back to the cache
assertThat(TestBundleProcessor.resetCnt, equalTo(1));
// BundleProcessor is added back to the BundleProcessorCache
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().size(), equalTo(1));
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(1));
// Add a reset handler that throws to test discarding the bundle processor on reset failure.
Iterables.getOnlyElement(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L")).getResetFunctions().add(() -> {
throw new IllegalStateException("ResetFailed");
});
handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("999L").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
// BundleProcessor is discarded instead of being added back to the BundleProcessorCache
assertThat(handler.bundleProcessorCache.getCachedBundleProcessors().get("1L").size(), equalTo(0));
}
Aggregations