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'))