def test_inactive_bundle_processor_returns_empty_split_response(self):
bundle_processor = mock.MagicMock()
bundle_processor_cache = BundleProcessorCache(None, None, {})
bundle_processor_cache.activate('instruction_id')
worker = SdkWorker(bundle_processor_cache)
split_request = beam_fn_api_pb2.InstructionRequest(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest(
instruction_id='instruction_id'))
self.assertEqual(
worker.do_instruction(split_request),
beam_fn_api_pb2.InstructionResponse(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.
ProcessBundleSplitResponse()))
# Add a mock bundle processor as if it was running before it's released
bundle_processor_cache.active_bundle_processors['instruction_id'] = (
'descriptor_id', bundle_processor)
bundle_processor_cache.release('instruction_id')
self.assertEqual(
worker.do_instruction(split_request),
beam_fn_api_pb2.InstructionResponse(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.
ProcessBundleSplitResponse()))
def test_failed_bundle_processor_returns_failed_split_response(self):
bundle_processor = mock.MagicMock()
bundle_processor_cache = BundleProcessorCache(None, None, {})
bundle_processor_cache.activate('instruction_id')
worker = SdkWorker(bundle_processor_cache)
# Add a mock bundle processor as if it was running before it's discarded
bundle_processor_cache.active_bundle_processors['instruction_id'] = (
'descriptor_id', bundle_processor)
bundle_processor_cache.discard('instruction_id')
split_request = beam_fn_api_pb2.InstructionRequest(
instruction_id='split_instruction_id',
process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest(
instruction_id='instruction_id'))
hc.assert_that(
worker.do_instruction(split_request).error,
hc.contains_string(
'Bundle processing associated with instruction_id has failed'))