def test_tracer_custom_annotation(mocker, dummy_response, provider_stub): put_annotation_mock = mocker.MagicMock() annotation_key = "BookingId" annotation_value = "123456" provider = provider_stub(put_annotation_mock=put_annotation_mock) tracer = Tracer(provider=provider, service="booking") tracer.put_annotation(annotation_key, annotation_value) assert put_annotation_mock.call_count == 1 assert put_annotation_mock.call_args == mocker.call(key=annotation_key, value=annotation_value)
class Utils: def __init__(self): self.logger = Logger() self.tracer = Tracer() def set_trace_id(self, lambda_handler): def get_trace_id(raw_trace_id): for t in raw_trace_id.split(';'): if t.startswith('Root'): return t.split('=')[1] return None def decorate(event, context): if '_X_AMZN_TRACE_ID' in os.environ: if get_trace_id(os.environ['_X_AMZN_TRACE_ID']) is not None: os.environ['trace_id'] = get_trace_id( os.environ['_X_AMZN_TRACE_ID']) if 'trace_id' not in os.environ: START_TIME = time.time() HEX = hex(int(START_TIME))[2:] os.environ['trace_id'] = "0-{}-{}".format( HEX, str(binascii.hexlify(urandom(12)), 'utf-8')) self.logger.structure_logs(append=True, trace_id=os.environ['trace_id']) self.logger.info(os.environ.get('trace_id')) response = lambda_handler(event, context) return response return decorate def api_gateway_response(self, lambda_handler): def decorate(event, context): response = lambda_handler(event, context) if isinstance(response, tuple): body = response[0] status = response[1] else: body = response status = 200 if isinstance(body, Exception): if status == 200: status = 400 self.logger.exception(body) self.tracer.put_annotation("ErrorCode", body.__class__.__name__) self.tracer.put_annotation("ErrorMessage", "{}".format(body)) res = { "isBase64Encoded": False, "statusCode": status, 'headers': { "Access-Control-Allow-Origin": "*", "Content-Type": "application/json" }, 'body': json.dumps({ "Code": body.__class__.__name__, "Message": "{}".format(body), "TraceId": os.environ.get('trace_id') }) } else: if body is not None and isinstance(body, str) is False: body = json.dumps(body) res = { 'isBase64Encoded': False, 'statusCode': status, 'headers': { "Access-Control-Allow-Origin": "*", "Content-Type": "application/json" } } if body is not None: res['body'] = body return res return decorate