def test_context(self): try: with statesampler.instruction_id('A'): tracker = statesampler.for_test() with tracker.scoped_state(NameContext('name', 'tid'), 'stage'): _LOGGER.info('message a') with statesampler.instruction_id('B'): _LOGGER.info('message b') _LOGGER.info('message c') self.fn_log_handler.close() a, b, c = sum([ list(logs.log_entries) for logs in self.test_logging_service.log_records_received ], []) self.assertEqual(a.instruction_id, 'A') self.assertEqual(b.instruction_id, 'B') self.assertEqual(c.instruction_id, '') self.assertEqual(a.transform_id, 'tid') self.assertEqual(b.transform_id, '') self.assertEqual(c.transform_id, '') finally: statesampler.set_current_tracker(None)
def _execute(self, task, request): with statesampler.instruction_id(request.instruction_id): try: response = task() except Exception: # pylint: disable=broad-except traceback_string = traceback.format_exc() print(traceback_string, file=sys.stderr) _LOGGER.error( 'Error processing instruction %s. Original traceback is\n%s\n', request.instruction_id, traceback_string) response = beam_fn_api_pb2.InstructionResponse( instruction_id=request.instruction_id, error=traceback_string) self._responses.put(response)