Skip to content

Commit ccba5ef

Browse files
dreamorosidreamorosi
committed
Completed captureLambdaHandler middleware
1 parent df2fc25 commit ccba5ef

File tree

4 files changed

+152
-47
lines changed

4 files changed

+152
-47
lines changed

‎packages/tracing/src/Tracer.ts

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -242,8 +242,7 @@ class Tracer implements TracerInterface {
242242
this.addResponseAsMetadata(result, context.functionName);
243243
} catch (error) {
244244
this.addErrorAsMetadata(error as Error);
245-
// TODO: should this error be thrown?? If thrown we get a ERR_UNHANDLED_REJECTION. If not aren't we are basically catching a Customer error?
246-
// throw error;
245+
throw error;
247246
} finally {
248247
subsegment?.close();
249248
}

‎packages/tracing/src/middleware/middy.ts

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,41 +3,45 @@ import { Subsegment } from 'aws-xray-sdk-core';
33
import { Tracer } from '../Tracer';
44

55
const captureLambdaHandler = (target: Tracer): middy.MiddlewareObj => {
6-
let subsegment: Subsegment | undefined;
76
const captureLambdaHandlerBefore = async (request: middy.Request): Promise<void> => {
87
if (target.isTracingEnabled()) {
9-
subsegment = new Subsegment(`## ${request.context.functionName}`);
8+
const subsegment = new Subsegment(`## ${request.context.functionName}`);
109
target.setSegment(subsegment);
1110

12-
if (Tracer.coldStart) {
11+
if (Tracer.isColdStart()) {
1312
target.putAnnotation('ColdStart', true);
1413
}
1514
}
1615
};
1716

1817
const captureLambdaHandlerAfter = async (request: middy.Request): Promise<void> => {
1918
if (target.isTracingEnabled()) {
20-
if (request.error) {
21-
if (target.isCaptureErrorEnabled() === false) {
22-
subsegment?.addErrorFlag();
23-
} else {
24-
subsegment?.addError(request.error, false);
25-
}
26-
// TODO: should this error be thrown??
27-
// throw error;
28-
} else {
29-
if (request.response !== undefined && target.isCaptureResponseEnabled() === true) {
30-
target.putMetadata(`${request.context.functionName} response`, request.response);
31-
}
19+
const subsegment = target.getSegment();
20+
if (request.response !== undefined && target.isCaptureResponseEnabled() === true) {
21+
target.putMetadata(`${request.context.functionName} response`, request.response);
3222
}
3323

3424
subsegment?.close();
3525
}
3626
};
3727

28+
const captureLambdaHandlerError = async (request: middy.Request): Promise<void> => {
29+
if (target.isTracingEnabled()) {
30+
const subsegment = target.getSegment();
31+
if (target.isCaptureErrorEnabled() === false) {
32+
subsegment?.addErrorFlag();
33+
} else {
34+
subsegment?.addError(request.error as Error, false);
35+
}
36+
// TODO: should this error be thrown?? I.e. should we stop the event flow & return?
37+
// throw error;
38+
}
39+
};
40+
3841
return {
3942
before: captureLambdaHandlerBefore,
4043
after: captureLambdaHandlerAfter,
44+
onError: captureLambdaHandlerError
4145
};
4246
};
4347

‎packages/tracing/tests/unit/Tracer.test.ts

Lines changed: 25 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@ import { context as dummyContext } from '../../../../tests/resources/contexts/he
22
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
33
// @ts-ignore
44
import * as dummyEvent from '../../../../tests/resources/events/custom/hello-world.json';
5-
// import { captureLambdaHandler } from '../../src/middleware/middy';
65
import { LambdaInterface } from '../../examples/utils/lambda';
7-
// import middy from '@middy/core';
86
import { Tracer } from '../../src';
97
import { Callback, Context } from 'aws-lambda/handler';
108
import { Segment, setContextMissingStrategy, Subsegment } from 'aws-xray-sdk-core';
@@ -454,18 +452,21 @@ describe('Class: Tracer', () => {
454452
}
455453

456454
// Act
457-
await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
458-
459-
// Assess
460-
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
461-
expect(newSubsegment).toEqual(expect.objectContaining({
462-
name: '## foo-bar-function',
463-
}));
464-
expect('cause' in newSubsegment).toBe(false);
465-
expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
466-
expect(addErrorSpy).toHaveBeenCalledTimes(0);
455+
try {
456+
await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
457+
} catch (error) {
458+
// Assess
459+
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
460+
expect(newSubsegment).toEqual(expect.objectContaining({
461+
name: '## foo-bar-function',
462+
}));
463+
expect('cause' in newSubsegment).toBe(false);
464+
expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
465+
expect(addErrorSpy).toHaveBeenCalledTimes(0);
466+
}
467467

468468
delete process.env.POWERTOOLS_TRACER_CAPTURE_ERROR;
469+
469470
});
470471

471472
test('when used as decorator and with standard config, it captures the exception correctly', async () => {
@@ -490,16 +491,18 @@ describe('Class: Tracer', () => {
490491
}
491492

492493
// Act
493-
await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
494-
495-
// Assess
496-
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
497-
expect(newSubsegment).toEqual(expect.objectContaining({
498-
name: '## foo-bar-function',
499-
}));
500-
expect('cause' in newSubsegment).toBe(true);
501-
expect(addErrorSpy).toHaveBeenCalledTimes(1);
502-
expect(addErrorSpy).toHaveBeenCalledWith(new Error('Exception thrown!'), false);
494+
try {
495+
await new Lambda().handler(dummyEvent, dummyContext, () => console.log('Lambda invoked!'));
496+
} catch (error) {
497+
// Assess
498+
expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(1);
499+
expect(newSubsegment).toEqual(expect.objectContaining({
500+
name: '## foo-bar-function',
501+
}));
502+
expect('cause' in newSubsegment).toBe(true);
503+
expect(addErrorSpy).toHaveBeenCalledTimes(1);
504+
expect(addErrorSpy).toHaveBeenCalledWith(new Error('Exception thrown!'), false);
505+
}
503506

504507
});
505508

‎packages/tracing/tests/unit/middy.test.ts

Lines changed: 107 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,32 @@ describe('Middy middlewares', () => {
6262

6363
});
6464

65-
test('when used as decorator while POWERTOOLS_TRACER_CAPTURE_RESPONSE is set to false, it does not capture the response as metadata', async () => {
65+
test('when used while tracing is disabled, even if the handler throws an error, it does nothing', async () => {
66+
67+
// Prepare
68+
const tracer: Tracer = new Tracer({ enabled: false });
69+
const setSegmentSpy = jest.spyOn(tracer.provider, 'setSegment').mockImplementation();
70+
const getSegmentSpy = jest.spyOn(tracer.provider, 'getSegment')
71+
.mockImplementationOnce(() => new Segment('facade', process.env._X_AMZN_TRACE_ID || null))
72+
.mockImplementationOnce(() => new Subsegment('## foo-bar-function'));
73+
const lambdaHandler: Handler = async (_event: unknown, _context: Context) => {
74+
throw new Error('Exception thrown!');
75+
};
76+
const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer));
77+
const context = Object.assign({}, mockContext);
78+
79+
// Act
80+
try {
81+
await handler({}, context, () => console.log('Lambda invoked!'));
82+
} catch (error) {
83+
// Assess
84+
expect(setSegmentSpy).toHaveBeenCalledTimes(0);
85+
expect(getSegmentSpy).toHaveBeenCalledTimes(0);
86+
}
87+
88+
});
89+
90+
test('when used while POWERTOOLS_TRACER_CAPTURE_RESPONSE is set to false, it does not capture the response as metadata', async () => {
6691

6792
// Prepare
6893
process.env.POWERTOOLS_TRACER_CAPTURE_RESPONSE = 'false';
@@ -87,7 +112,7 @@ describe('Middy middlewares', () => {
87112

88113
});
89114

90-
test('when used as decorator and with standard config, it captures the response as metadata', async () => {
115+
test('when used with standard config, it captures the response as metadata', async () => {
91116

92117
// Prepare
93118
const tracer: Tracer = new Tracer();
@@ -122,7 +147,7 @@ describe('Middy middlewares', () => {
122147

123148
});
124149

125-
test('when used as decorator while POWERTOOLS_TRACER_CAPTURE_ERROR is set to false, it does not capture the exceptions', async () => {
150+
test('when used while POWERTOOLS_TRACER_CAPTURE_ERROR is set to false, it does not capture the exceptions', async () => {
126151

127152
// Prepare
128153
process.env.POWERTOOLS_TRACER_CAPTURE_ERROR = 'false';
@@ -140,20 +165,94 @@ describe('Middy middlewares', () => {
140165
const context = Object.assign({}, mockContext);
141166

142167
// Act
143-
await handler({}, context, () => console.log('Lambda invoked!'));
168+
try {
169+
await handler({}, context, () => console.log('Lambda invoked!'));
170+
} catch (error) {
171+
// Assess
172+
expect(setSegmentSpy).toHaveBeenCalledTimes(1);
173+
expect(setSegmentSpy).toHaveBeenCalledWith(expect.objectContaining({
174+
name: '## foo-bar-function',
175+
}));
176+
expect('cause' in newSubsegment).toBe(false);
177+
expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
178+
expect(addErrorSpy).toHaveBeenCalledTimes(0);
179+
}
180+
181+
delete process.env.POWERTOOLS_TRACER_CAPTURE_ERROR;
182+
183+
});
184+
185+
});
144186

187+
test('when used with standard config, it captures the exception correctly', async () => {
188+
189+
// Prepare
190+
const tracer: Tracer = new Tracer();
191+
const newSubsegment: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
192+
const setSegmentSpy = jest.spyOn(tracer.provider, 'setSegment').mockImplementation();
193+
jest.spyOn(tracer.provider, 'getSegment').mockImplementation(() => newSubsegment);
194+
setContextMissingStrategy(() => null);
195+
const addErrorSpy = jest.spyOn(newSubsegment, 'addError');
196+
const lambdaHandler: Handler = async (_event: unknown, _context: Context) => {
197+
throw new Error('Exception thrown!');
198+
};
199+
const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer));
200+
const context = Object.assign({}, mockContext);
201+
202+
// Act
203+
try {
204+
await handler({}, context, () => console.log('Lambda invoked!'));
205+
} catch (error) {
145206
// Assess
146207
expect(setSegmentSpy).toHaveBeenCalledTimes(1);
147208
expect(setSegmentSpy).toHaveBeenCalledWith(expect.objectContaining({
148209
name: '## foo-bar-function',
149210
}));
150-
expect('cause' in newSubsegment).toBe(false);
151-
expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
152-
expect(addErrorSpy).toHaveBeenCalledTimes(0);
211+
expect('cause' in newSubsegment).toBe(true);
212+
expect(addErrorSpy).toHaveBeenCalledTimes(1);
213+
expect(addErrorSpy).toHaveBeenCalledWith(new Error('Exception thrown!'), false);
214+
}
153215

154-
delete process.env.POWERTOOLS_TRACER_CAPTURE_ERROR;
216+
});
155217

218+
test('when used with standard config, it annotates ColdStart correctly', async () => {
219+
220+
// Prepare
221+
const tracer: Tracer = new Tracer();
222+
const newSubsegmentFirstInvocation: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
223+
const newSubsegmentSecondInvocation: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
224+
const setSegmentSpy = jest.spyOn(tracer.provider, 'setSegment').mockImplementation();
225+
jest.spyOn(tracer.provider, 'getSegment')
226+
.mockImplementationOnce(() => newSubsegmentFirstInvocation)
227+
.mockImplementation(() => newSubsegmentSecondInvocation);
228+
setContextMissingStrategy(() => null);
229+
const addAnnotationSpy = jest.spyOn(tracer, 'putAnnotation');
230+
const lambdaHandler: Handler = async (_event: unknown, _context: Context) => ({
231+
foo: 'bar'
156232
});
233+
const handler = middy(lambdaHandler).use(captureLambdaHandler(tracer));
234+
const context = Object.assign({}, mockContext);
235+
236+
// Act
237+
await handler({}, context, () => console.log('Lambda invoked!'));
238+
await handler({}, context, () => console.log('Lambda invoked!'));
239+
240+
// Assess
241+
expect(setSegmentSpy).toHaveBeenCalledTimes(2);
242+
expect(setSegmentSpy).toHaveBeenCalledWith(expect.objectContaining({
243+
name: '## foo-bar-function',
244+
}));
245+
expect(addAnnotationSpy).toHaveBeenCalledTimes(1);
246+
expect(addAnnotationSpy).toHaveBeenCalledWith('ColdStart', true);
247+
expect(newSubsegmentFirstInvocation).toEqual(expect.objectContaining({
248+
name: '## foo-bar-function',
249+
annotations: {
250+
'ColdStart': true,
251+
}
252+
}));
253+
expect(newSubsegmentSecondInvocation).toEqual(expect.objectContaining({
254+
name: '## foo-bar-function'
255+
}));
157256

158257
});
159258

0 commit comments

Comments
 (0)