Skip to content

Commit 640a431

Browse files
dreamorosidreamorosi
committed
Added captureAWS feature to e2e
1 parent 196db23 commit 640a431

File tree

5 files changed

+182
-39
lines changed

5 files changed

+182
-39
lines changed

‎packages/tracing/tests/e2e/tracer.test.Decorator.ts

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
import { Tracer } from '../../src';
22
import { Callback, Context } from 'aws-lambda';
3+
import { STSClient, GetCallerIdentityCommand } from '@aws-sdk/client-sts';
4+
// eslint-disable-next-line @typescript-eslint/no-var-requires
5+
let AWS = require('aws-sdk');
36

47
const serviceName = process.env.EXPECTED_SERVICE_NAME ?? 'MyFunctionWithStandardHandler';
58
const customAnnotationKey = process.env.EXPECTED_CUSTOM_ANNOTATION_KEY ?? 'myAnnotation';
@@ -11,9 +14,24 @@ const customErrorMessage = process.env.EXPECTED_CUSTOM_ERROR_MESSAGE ?? 'An erro
1114

1215
interface CustomEvent {
1316
throw: boolean
17+
sdkV2: string
18+
invocation: number
1419
}
1520

21+
// Function that refreshes imports to ensure that we are instrumenting only one version of the AWS SDK v2 at a time.
22+
const refreshAWSSDKImport = (): void => {
23+
// Clean up the require cache to ensure we're using a newly imported version of the AWS SDK v2
24+
for (const key in require.cache) {
25+
if (key.indexOf('/aws-sdk/') !== -1) {
26+
delete require.cache[key];
27+
}
28+
}
29+
// eslint-disable-next-line @typescript-eslint/no-var-requires
30+
AWS = require('aws-sdk');
31+
};
32+
1633
const tracer = new Tracer({ serviceName: serviceName });
34+
const stsv3 = tracer.captureAWSv3Client(new STSClient({}));
1735

1836
export class MyFunctionWithDecorator {
1937
@tracer.captureLambdaHandler()
@@ -23,7 +41,35 @@ export class MyFunctionWithDecorator {
2341
tracer.putAnnotation(customAnnotationKey, customAnnotationValue);
2442
tracer.putMetadata(customMetadataKey, customMetadataValue);
2543

26-
let res;
44+
let stsv2;
45+
refreshAWSSDKImport();
46+
if (event.sdkV2 === 'client') {
47+
stsv2 = tracer.captureAWSClient(new AWS.STS());
48+
} else if (event.sdkV2 === 'all') {
49+
AWS = tracer.captureAWS(AWS);
50+
stsv2 = new AWS.STS();
51+
}
52+
53+
return Promise.all([
54+
stsv2.getCallerIdentity().promise(),
55+
stsv3.send(new GetCallerIdentityCommand({})),
56+
new Promise((resolve, reject) => {
57+
setTimeout(() => {
58+
const res = this.myMethod();
59+
if (event.throw) {
60+
reject(new Error(customErrorMessage));
61+
} else {
62+
resolve(res);
63+
}
64+
}, 2000); // We need to wait for to make sure previous calls are finished
65+
})
66+
])
67+
.then(([ _stsv2Res, _stsv3Res, promiseRes ]) => promiseRes)
68+
.catch((err) => {
69+
throw err;
70+
});
71+
72+
/* let res;
2773
try {
2874
res = this.myMethod();
2975
if (event.throw) {
@@ -33,7 +79,7 @@ export class MyFunctionWithDecorator {
3379
throw err;
3480
}
3581
36-
return res;
82+
return res; */
3783
}
3884

3985
@tracer.captureMethod()

‎packages/tracing/tests/e2e/tracer.test.Manual.ts

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import { Tracer } from '../../src';
22
import { Context } from 'aws-lambda';
33
import { STSClient, GetCallerIdentityCommand } from '@aws-sdk/client-sts';
4-
import { STS } from 'aws-sdk';
4+
// eslint-disable-next-line @typescript-eslint/no-var-requires
5+
let AWS = require('aws-sdk');
56

67
const serviceName = process.env.EXPECTED_SERVICE_NAME ?? 'MyFunctionWithStandardHandler';
78
const customAnnotationKey = process.env.EXPECTED_CUSTOM_ANNOTATION_KEY ?? 'myAnnotation';
@@ -13,10 +14,23 @@ const customErrorMessage = process.env.EXPECTED_CUSTOM_ERROR_MESSAGE ?? 'An erro
1314

1415
interface CustomEvent {
1516
throw: boolean
17+
sdkV2: string
18+
invocation: number
1619
}
1720

21+
// Function that refreshes imports to ensure that we are instrumenting only one version of the AWS SDK v2 at a time.
22+
const refreshAWSSDKImport = (): void => {
23+
// Clean up the require cache to ensure we're using a newly imported version of the AWS SDK v2
24+
for (const key in require.cache) {
25+
if (key.indexOf('/aws-sdk/') !== -1) {
26+
delete require.cache[key];
27+
}
28+
}
29+
// eslint-disable-next-line @typescript-eslint/no-var-requires
30+
AWS = require('aws-sdk');
31+
};
32+
1833
const tracer = new Tracer({ serviceName: serviceName });
19-
const stsv2 = tracer.captureAWSClient(new STS({}));
2034
const stsv3 = tracer.captureAWSv3Client(new STSClient({}));
2135

2236
export const handler = async (event: CustomEvent, _context: Context): Promise<void> => {
@@ -27,9 +41,18 @@ export const handler = async (event: CustomEvent, _context: Context): Promise<vo
2741
tracer.annotateColdStart();
2842
tracer.addServiceNameAnnotation();
2943

44+
tracer.putAnnotation('invocation', event.invocation);
3045
tracer.putAnnotation(customAnnotationKey, customAnnotationValue);
3146
tracer.putMetadata(customMetadataKey, customMetadataValue);
3247

48+
let stsv2;
49+
refreshAWSSDKImport();
50+
if (event.sdkV2 === 'client') {
51+
stsv2 = tracer.captureAWSClient(new AWS.STS());
52+
} else if (event.sdkV2 === 'all') {
53+
AWS = tracer.captureAWS(AWS);
54+
stsv2 = new AWS.STS();
55+
}
3356
try {
3457
await stsv2.getCallerIdentity().promise();
3558
} catch (err) {

‎packages/tracing/tests/e2e/tracer.test.Middleware.ts

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
import middy from '@middy/core';
22
import { captureLambdaHandler, Tracer } from '../../src';
33
import { Context } from 'aws-lambda';
4+
import { STSClient, GetCallerIdentityCommand } from '@aws-sdk/client-sts';
5+
// eslint-disable-next-line @typescript-eslint/no-var-requires
6+
let AWS = require('aws-sdk');
47

58
const serviceName = process.env.EXPECTED_SERVICE_NAME ?? 'MyFunctionWithStandardHandler';
69
const customAnnotationKey = process.env.EXPECTED_CUSTOM_ANNOTATION_KEY ?? 'myAnnotation';
@@ -12,14 +15,50 @@ const customErrorMessage = process.env.EXPECTED_CUSTOM_ERROR_MESSAGE ?? 'An erro
1215

1316
interface CustomEvent {
1417
throw: boolean
18+
sdkV2: string
19+
invocation: number
1520
}
1621

22+
// Function that refreshes imports to ensure that we are instrumenting only one version of the AWS SDK v2 at a time.
23+
const refreshAWSSDKImport = (): void => {
24+
// Clean up the require cache to ensure we're using a newly imported version of the AWS SDK v2
25+
for (const key in require.cache) {
26+
if (key.indexOf('/aws-sdk/') !== -1) {
27+
delete require.cache[key];
28+
}
29+
}
30+
// eslint-disable-next-line @typescript-eslint/no-var-requires
31+
AWS = require('aws-sdk');
32+
};
33+
1734
const tracer = new Tracer({ serviceName: serviceName });
35+
const stsv3 = tracer.captureAWSv3Client(new STSClient({}));
1836

1937
export const handler = middy(async (event: CustomEvent, _context: Context): Promise<void> => {
38+
tracer.putAnnotation('invocation', event.invocation);
2039
tracer.putAnnotation(customAnnotationKey, customAnnotationValue);
2140
tracer.putMetadata(customMetadataKey, customMetadataValue);
2241

42+
let stsv2;
43+
refreshAWSSDKImport();
44+
if (event.sdkV2 === 'client') {
45+
stsv2 = tracer.captureAWSClient(new AWS.STS());
46+
} else if (event.sdkV2 === 'all') {
47+
AWS = tracer.captureAWS(AWS);
48+
stsv2 = new AWS.STS();
49+
}
50+
try {
51+
await stsv2.getCallerIdentity().promise();
52+
} catch (err) {
53+
console.error(err);
54+
}
55+
56+
try {
57+
await stsv3.send(new GetCallerIdentityCommand({}));
58+
} catch (err) {
59+
console.error(err);
60+
}
61+
2362
let res;
2463
try {
2564
res = customResponseValue;

‎packages/tracing/tests/e2e/tracer.test.ts

Lines changed: 65 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { SdkProvider } from 'aws-cdk/lib/api/aws-auth';
1212
import { CloudFormationDeployments } from 'aws-cdk/lib/api/cloudformation-deployments';
1313
import * as AWS from 'aws-sdk';
1414
import { getTraces, getInvocationSubsegment } from '../helpers/tracesUtils';
15+
import type { ParsedDocument } from '../helpers/tracesUtils';
1516

1617
const xray = new AWS.XRay();
1718
const lambdaClient = new AWS.Lambda();
@@ -88,7 +89,9 @@ describe('Tracer integration tests', () => {
8889
await lambdaClient.invoke({
8990
FunctionName: functionName,
9091
Payload: JSON.stringify({
91-
throw: i === invocations ? true : false, // only last invocation should throw
92+
throw: i === invocations - 1 ? true : false, // only last invocation should throw
93+
sdkV2: i === 1 ? 'all' : 'client', // only second invocation should use captureAll
94+
invocation: i + 1, // Pass invocation number for easier debugging
9295
}),
9396
}).promise();
9497
}
@@ -197,7 +200,7 @@ describe('Tracer integration tests', () => {
197200

198201
for (let i = 0; i < invocations; i++) {
199202
// Assert that the trace has the expected amount of segments
200-
expect(sortedTraces[i].Segments.length).toBe(2);
203+
expect(sortedTraces[i].Segments.length).toBe(4);
201204

202205
const invocationSubsegment = getInvocationSubsegment(sortedTraces[i]);
203206

@@ -206,35 +209,46 @@ describe('Tracer integration tests', () => {
206209
const handlerSubsegment = invocationSubsegment?.subsegments[0];
207210
// Assert that the subsegment name is the expected one
208211
expect(handlerSubsegment.name).toBe('## index.handler');
209-
// Assert that there're no subsegments
210-
expect(handlerSubsegment.hasOwnProperty('subsegments')).toBe(false);
211-
212-
const { annotations, metadata } = handlerSubsegment;
212+
if (handlerSubsegment?.subsegments !== undefined) {
213+
// Assert that there're two subsegments
214+
expect(handlerSubsegment?.subsegments?.length).toBe(2);
213215

214-
if (annotations !== undefined && metadata !== undefined) {
215-
// Assert that the annotations are as expected
216-
expect(annotations['ColdStart']).toEqual(true ? i === 0 : false);
217-
expect(annotations['Service']).toEqual(expectedServiceName);
218-
expect(annotations[expectedCustomAnnotationKey]).toEqual(expectedCustomAnnotationValue);
219-
// Assert that the metadata object is as expected
220-
expect(metadata[expectedServiceName][expectedCustomMetadataKey])
221-
.toEqual(expectedCustomMetadataValue);
216+
const [ AWSSDKSubsegment1, AWSSDKSubsegment2 ] = handlerSubsegment?.subsegments;
217+
// Assert that the subsegment names is the expected ones
218+
expect(AWSSDKSubsegment1.name).toBe('STS');
219+
expect(AWSSDKSubsegment2.name).toBe('STS');
222220

223-
if (i === invocations - 1) {
224-
// Assert that the subsegment has the expected fault
225-
expect(invocationSubsegment.error).toBe(true);
226-
expect(handlerSubsegment.fault).toBe(true);
227-
expect(handlerSubsegment.hasOwnProperty('cause')).toBe(true);
228-
expect(handlerSubsegment.cause?.exceptions[0].message).toBe(expectedCustomErrorMessage);
221+
const { annotations, metadata } = handlerSubsegment;
222+
223+
if (annotations !== undefined && metadata !== undefined) {
224+
// Assert that the annotations are as expected
225+
expect(annotations['ColdStart']).toEqual(true ? i === 0 : false);
226+
expect(annotations['Service']).toEqual(expectedServiceName);
227+
expect(annotations[expectedCustomAnnotationKey]).toEqual(expectedCustomAnnotationValue);
228+
// Assert that the metadata object is as expected
229+
expect(metadata[expectedServiceName][expectedCustomMetadataKey])
230+
.toEqual(expectedCustomMetadataValue);
231+
232+
if (i === invocations - 1) {
233+
// Assert that the subsegment has the expected fault
234+
expect(invocationSubsegment.error).toBe(true);
235+
expect(handlerSubsegment.fault).toBe(true);
236+
expect(handlerSubsegment.hasOwnProperty('cause')).toBe(true);
237+
expect(handlerSubsegment.cause?.exceptions[0].message).toBe(expectedCustomErrorMessage);
238+
} else {
239+
// Assert that the metadata object contains the response
240+
expect(metadata[expectedServiceName]['index.handler response'])
241+
.toEqual(expectedCustomResponseValue);
242+
}
229243
} else {
230-
// Assert that the metadata object contains the response
231-
expect(metadata[expectedServiceName]['index.handler response'])
232-
.toEqual(expectedCustomResponseValue);
244+
// Make test fail if there are no annotations or metadata
245+
expect('annotations !== undefined && metadata !== undefined')
246+
.toBe('annotations === undefined && metadata === undefined');
233247
}
234248
} else {
235-
// Make test fail if there are no annotations or metadata
236-
expect('annotations !== undefined && metadata !== undefined')
237-
.toBe('annotations === undefined && metadata === undefined');
249+
// Make test fail if the handlerSubsegment subsegment doesn't have any subsebment
250+
expect('handlerSubsegment?.subsegments !== undefined')
251+
.toBe('handlerSubsegment?.subsegments === undefined');
238252
}
239253
} else {
240254
// Make test fail if the Invocation subsegment doesn't have an handler subsebment
@@ -256,7 +270,7 @@ describe('Tracer integration tests', () => {
256270

257271
for (let i = 0; i < invocations; i++) {
258272
// Assert that the trace has the expected amount of segments
259-
expect(sortedTraces[i].Segments.length).toBe(2);
273+
expect(sortedTraces[i].Segments.length).toBe(4);
260274

261275
const invocationSubsegment = getInvocationSubsegment(sortedTraces[i]);
262276

@@ -266,13 +280,30 @@ describe('Tracer integration tests', () => {
266280
// Assert that the subsegment name is the expected one
267281
expect(handlerSubsegment.name).toBe('## index.handler');
268282
if (handlerSubsegment?.subsegments !== undefined) {
269-
// Assert that there is one subsegment
270-
expect(handlerSubsegment?.subsegments?.length).toBe(1);
271-
const methodSubsegment = handlerSubsegment?.subsegments[0];
272-
// Assert that the subsegment name is the expected one
273-
expect(methodSubsegment.name).toBe('### myMethod');
274-
275-
const { metadata } = methodSubsegment;
283+
// Assert that there're three subsegments
284+
expect(handlerSubsegment?.subsegments?.length).toBe(3);
285+
286+
// Sort the subsegments by name
287+
const stsSubsegments: ParsedDocument[] = [];
288+
const methodSubsegment: ParsedDocument[] = [];
289+
const otherSegments: ParsedDocument[] = [];
290+
handlerSubsegment?.subsegments.forEach(subsegment => {
291+
if (subsegment.name === 'STS') {
292+
stsSubsegments.push(subsegment);
293+
} else if (subsegment.name === '### myMethod') {
294+
methodSubsegment.push(subsegment);
295+
} else {
296+
otherSegments.push(subsegment);
297+
}
298+
});
299+
// Assert that there are exactly two subsegment with the name 'STS'
300+
expect(stsSubsegments.length).toBe(2);
301+
// Assert that there is exactly one subsegment with the name '### myMethod'
302+
expect(methodSubsegment.length).toBe(1);
303+
// Assert that there are exactly zero other subsegments
304+
expect(otherSegments.length).toBe(0);
305+
306+
const { metadata } = methodSubsegment[0];
276307

277308
if (metadata !== undefined) {
278309
// Assert that the metadata object is as expected

‎packages/tracing/tests/helpers/tracesUtils.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,5 +124,9 @@ const getInvocationSubsegment = (trace: ParsedTrace): ParsedDocument => {
124124
export {
125125
getTraces,
126126
getFunctionSegment,
127-
getInvocationSubsegment
127+
getInvocationSubsegment,
128+
};
129+
130+
export type {
131+
ParsedDocument,
128132
};

0 commit comments

Comments
 (0)