test: add middy options e2e

Author: misterjoshua

Diff for: packages/tracer/tests/e2e/allFeatures.middy.test.functionCode.ts

@@ -36,7 +36,7 @@ const refreshAWSSDKImport = (): void => {
 const tracer = new Tracer({ serviceName: serviceName });
 const dynamoDBv3 = tracer.captureAWSv3Client(new DynamoDBClient({}));
 
-export const handler = middy(async (event: CustomEvent, _context: Context): Promise<void> => {
+const testHandler = async (event: CustomEvent, _context: Context): Promise<void> => {
   tracer.putAnnotation('invocation', event.invocation);
   tracer.putAnnotation(customAnnotationKey, customAnnotationValue);
   tracer.putMetadata(customMetadataKey, customMetadataValue);
@@ -63,4 +63,8 @@ export const handler = middy(async (event: CustomEvent, _context: Context): Prom
   } catch (err) {
     throw err;
   }
-}).use(captureLambdaHandler(tracer));
+};
+
+export const handler = middy(testHandler).use(captureLambdaHandler(tracer));
+
+export const handlerWithNoCaptureResponseViaMiddlewareOption = middy(testHandler).use(captureLambdaHandler(tracer, { captureResponse: false }));

Diff for: packages/tracer/tests/e2e/allFeatures.middy.test.ts

@@ -78,6 +78,13 @@ const uuidFunction3 = v4();
 const functionNameWithTracerDisabled = generateUniqueName(RESOURCE_NAME_PREFIX, uuidFunction3, runtime, 'AllFeatures-Middy-TracerDisabled');
 const serviceNameWithTracerDisabled = functionNameWithNoCaptureErrorOrResponse; 
 
+/**
+ * Function #4 doesn't capture response
+ */
+const uuidFunction4 = v4();
+const functionNameWithNoCaptureResponseViaMiddlewareOption = generateUniqueName(RESOURCE_NAME_PREFIX, uuidFunction4, runtime, 'AllFeatures-Middy-NoCaptureResponse2');
+const serviceNameWithNoCaptureResponseViaMiddlewareOption = functionNameWithNoCaptureResponseViaMiddlewareOption; 
+
 const xray = new AWS.XRay();
 const invocations = 3;
 
@@ -149,13 +156,30 @@ describe(`Tracer E2E tests, all features with middy instantiation for runtime: $
     });
     ddbTable.grantWriteData(functionWithTracerDisabled);
 
+    const functionThatDoesNotCaptureResponseViaMiddlewareOption = createTracerTestFunction({
+      stack,
+      functionName: functionNameWithNoCaptureResponseViaMiddlewareOption,
+      entry,
+      handler: 'handlerWithNoCaptureResponseViaMiddlewareOption',
+      expectedServiceName: serviceNameWithNoCaptureResponseViaMiddlewareOption,
+      environmentParams: {
+        TEST_TABLE_NAME: ddbTableName,
+        POWERTOOLS_TRACER_CAPTURE_RESPONSE: 'true',
+        POWERTOOLS_TRACER_CAPTURE_ERROR: 'true',
+        POWERTOOLS_TRACE_ENABLED: 'true',
+      },
+      runtime
+    });
+    ddbTable.grantWriteData(functionThatDoesNotCaptureResponseViaMiddlewareOption);
+
     await deployStack(integTestApp, stack);
 
     // Act
     await Promise.all([
       invokeAllTestCases(functionNameWithAllFlagsEnabled),
       invokeAllTestCases(functionNameWithNoCaptureErrorOrResponse),
       invokeAllTestCases(functionNameWithTracerDisabled),
+      invokeAllTestCases(functionNameWithNoCaptureResponseViaMiddlewareOption),
     ]);
 
   }, SETUP_TIMEOUT);
@@ -299,6 +323,54 @@ describe(`Tracer E2E tests, all features with middy instantiation for runtime: $
 
   }, TEST_CASE_TIMEOUT);
 
+  it('should not capture response when the middleware\'s captureResponse is set to false', async () => {
+    
+    const tracesWithNoCaptureResponse = await getTraces(xray, startTime, await getFunctionArn(functionNameWithNoCaptureResponseViaMiddlewareOption), invocations, 5);
+    
+    expect(tracesWithNoCaptureResponse.length).toBe(invocations);
+
+    // Assess
+    for (let i = 0; i < invocations; i++) {
+      const trace = tracesWithNoCaptureResponse[i];
+
+      /**
+       * Expect the trace to have 5 segments:
+       * 1. Lambda Context (AWS::Lambda)
+       * 2. Lambda Function (AWS::Lambda::Function)
+       * 3. DynamoDB (AWS::DynamoDB)
+       * 4. DynamoDB Table (AWS::DynamoDB::Table)
+       * 5. Remote call (httpbin.org)
+       */
+      expect(trace.Segments.length).toBe(5);
+      const invocationSubsegment = getInvocationSubsegment(trace);
+      
+      /**
+       * Invocation subsegment should have a subsegment '## index.handlerWithNoCaptureResponseViaMiddlewareOption' (default behavior for PowerTool tracer)
+       * '## index.handlerWithNoCaptureResponseViaMiddlewareOption' subsegment should have 3 subsegments
+       * 1. DynamoDB (PutItem on the table)
+       * 2. DynamoDB (PutItem overhead)
+       * 3. httpbin.org (Remote call)
+       */
+      const handlerSubsegment = getFirstSubsegment(invocationSubsegment);
+      expect(handlerSubsegment.name).toBe('## index.handlerWithNoCaptureResponseViaMiddlewareOption');
+      expect(handlerSubsegment?.subsegments).toHaveLength(3);
+
+      if (!handlerSubsegment.subsegments) {
+        fail('"## index.handlerWithNoCaptureResponseViaMiddlewareOption" subsegment should have subsegments');
+      }
+      const subsegments = splitSegmentsByName(handlerSubsegment.subsegments, [ 'DynamoDB', 'httpbin.org' ]);
+      expect(subsegments.get('DynamoDB')?.length).toBe(2);
+      expect(subsegments.get('httpbin.org')?.length).toBe(1);
+      expect(subsegments.get('other')?.length).toBe(0);
+      
+      const shouldThrowAnError = (i === (invocations - 1));
+      if (shouldThrowAnError) {
+        assertErrorAndFault(invocationSubsegment, expectedCustomErrorMessage);
+      }
+    }
+
+  }, TEST_CASE_TIMEOUT);
+
   it('should not capture any custom traces when disabled', async () => {
     const expectedNoOfTraces = 2;
     const tracesWithTracerDisabled = await getTraces(xray, startTime, await getFunctionArn(functionNameWithTracerDisabled), invocations, expectedNoOfTraces);

Diff for: packages/tracer/tests/helpers/tracesUtils.ts

@@ -79,6 +79,7 @@ export interface ParsedTrace {
 interface TracerTestFunctionParams { 
   stack: Stack
   functionName: string
+  handler?: string
   entry: string
   expectedServiceName: string
   environmentParams: { [key: string]: string }
@@ -237,7 +238,7 @@ const createTracerTestFunction = (params: TracerTestFunctionParams): NodejsFunct
   const func = new NodejsFunction(stack, functionName, {
     entry: entry,
     functionName: functionName,
-    handler: 'handler',
+    handler: params.handler ?? 'handler',
     tracing: Tracing.ACTIVE,
     architecture: Architecture.X86_64,
     memorySize: 256, // Default value (128) will take too long to process