test(tracer): streamline e2e test structure

packages/testing/src/types.ts

@@ -108,8 +108,13 @@ type XRayTraceDocumentParsed = {
     request_id: string;
   };
   http?: {
-    response: {
+    request: {
+      url: string;
+      method: string;
+    };
+    response?: {
       status: number;
+      content_length?: number;
     };
   };
   origin?: string;
@@ -142,6 +147,7 @@ type XRayTraceDocumentParsed = {
     message: string;
   };
   error?: boolean;
+  namespace?: string;
 };
 
 type XRaySegmentParsed = {
@@ -163,6 +169,10 @@ type GetXRayTraceDetailsOptions = {
    * The expected number of segments in each trace
    */
   expectedSegmentsCount: number;
+  /**
+   * The name of the function that the trace is expected to be associated with
+   */
+  functionName: string;
 };
 
 /**

packages/testing/src/xray-traces-utils.ts

@@ -83,18 +83,66 @@ const retriableGetTraceIds = (options: GetXRayTraceIdsOptions) =>
             endTime.getTime() / 1000
           )} --filter-expression 'resource.arn ENDSWITH ":function:${options.resourceName}"'`
         );
+
+        throw new Error(
+          `Failed to get trace IDs after ${retryOptions.retries} retries`,
+          { cause: error }
+        );
       }
       retry(error);
     }
-  });
+  }, retryOptions);
+
+/**
+ * Find the main Powertools subsegment in the trace
+ *
+ * A main Powertools subsegment is identified by the `## index.` suffix. Depending on the
+ * runtime, it may also be identified by the `Invocation` name.
+ *
+ * @param trace - The trace to find the main Powertools subsegment
+ * @param functionName - The function name to find the main Powertools subsegment
+ */
+const findMainPowertoolsSubsegment = (
+  trace: XRayTraceDocumentParsed,
+  functionName: string
+) => {
+  const maybePowertoolsSubsegment = trace.subsegments?.find(
+    (subsegment) =>
+      subsegment.name.startsWith('## index.') ||
+      subsegment.name === 'Invocation'
+  );
+
+  if (!maybePowertoolsSubsegment) {
+    throw new Error(`Main subsegment not found for ${functionName} segment`);
+  }
+
+  if (maybePowertoolsSubsegment.name === 'Invocation') {
+    const powertoolsSubsegment = maybePowertoolsSubsegment.subsegments?.find(
+      (subsegment) => subsegment.name.startsWith('## index.')
+    );
+
+    if (!powertoolsSubsegment) {
+      throw new Error(`Main subsegment not found for ${functionName} segment`);
+    }
+
+    return powertoolsSubsegment;
+  }
+
+  return maybePowertoolsSubsegment;
+};
 
 /**
  * Parse and sort the trace segments by start time
  *
  * @param trace - The trace to parse and sort
  * @param expectedSegmentsCount - The expected segments count for the trace
+ * @param functionName - The function name to find the main Powertools subsegment
  */
-const parseAndSortTrace = (trace: Trace, expectedSegmentsCount: number) => {
+const parseAndSortTrace = (
+  trace: Trace,
+  expectedSegmentsCount: number,
+  functionName: string
+) => {
   const { Id: id, Segments: segments } = trace;
   if (segments === undefined || segments.length !== expectedSegmentsCount) {
     throw new Error(
@@ -111,9 +159,14 @@ const parseAndSortTrace = (trace: Trace, expectedSegmentsCount: number) => {
       );
     }
 
+    const parsedDocument = JSON.parse(Document) as XRayTraceDocumentParsed;
+    if (parsedDocument.origin === 'AWS::Lambda::Function') {
+      findMainPowertoolsSubsegment(parsedDocument, functionName);
+    }
+
     parsedSegments.push({
       Id,
-      Document: JSON.parse(Document) as XRayTraceDocumentParsed,
+      Document: parsedDocument,
     });
   }
 
@@ -136,15 +189,14 @@ const parseAndSortTrace = (trace: Trace, expectedSegmentsCount: number) => {
 const getTraceDetails = async (
   options: GetXRayTraceDetailsOptions
 ): Promise<XRayTraceParsed[]> => {
-  const { traceIds, expectedSegmentsCount } = options;
+  const { traceIds, expectedSegmentsCount, functionName } = options;
   const response = await xrayClient.send(
     new BatchGetTracesCommand({
       TraceIds: traceIds,
     })
   );
 
-  const traces = response.Traces;
-
+  const { Traces: traces } = response;
   if (traces === undefined || traces.length !== traceIds.length) {
     throw new Error(
       `Expected ${traceIds.length} traces, got ${traces ? traces.length : 0}`
@@ -153,7 +205,9 @@ const getTraceDetails = async (
 
   const parsedAndSortedTraces: XRayTraceParsed[] = [];
   for (const trace of traces) {
-    parsedAndSortedTraces.push(parseAndSortTrace(trace, expectedSegmentsCount));
+    parsedAndSortedTraces.push(
+      parseAndSortTrace(trace, expectedSegmentsCount, functionName)
+    );
   }
 
   return parsedAndSortedTraces.sort(
@@ -168,16 +222,28 @@ const getTraceDetails = async (
  * @param options - The options to get trace details, including the trace IDs and expected segments count
  */
 const retriableGetTraceDetails = (options: GetXRayTraceDetailsOptions) =>
-  promiseRetry(async (retry) => {
+  promiseRetry(async (retry, attempt) => {
     try {
       return await getTraceDetails(options);
     } catch (error) {
+      if (attempt === retryOptions.retries) {
+        console.log(
+          `Manual query: aws xray batch-get-traces --trace-ids ${
+            options.traceIds
+          }`
+        );
+
+        throw new Error(
+          `Failed to get trace details after ${retryOptions.retries} retries`,
+          { cause: error }
+        );
+      }
       retry(error);
     }
-  });
+  }, retryOptions);
 
 /**
- * Find the main function segment in the trace identified by the `## index.` suffix
+ * Find the main function segment within the `AWS::Lambda::Function` segment
  */
 const findPowertoolsFunctionSegment = (
   trace: XRayTraceParsed,
@@ -194,30 +260,7 @@ const findPowertoolsFunctionSegment = (
   }
 
   const document = functionSegment.Document;
-
-  const maybePowertoolsSubsegment = document.subsegments?.find(
-    (subsegment) =>
-      subsegment.name.startsWith('## index.') ||
-      subsegment.name === 'Invocation'
-  );
-
-  if (!maybePowertoolsSubsegment) {
-    throw new Error(`Main subsegment not found for ${functionName} segment`);
-  }
-
-  if (maybePowertoolsSubsegment.name === 'Invocation') {
-    const powertoolsSubsegment = maybePowertoolsSubsegment.subsegments?.find(
-      (subsegment) => subsegment.name.startsWith('## index.')
-    );
-
-    if (!powertoolsSubsegment) {
-      throw new Error(`Main subsegment not found for ${functionName} segment`);
-    }
-
-    return powertoolsSubsegment;
-  }
-
-  return maybePowertoolsSubsegment;
+  return findMainPowertoolsSubsegment(document, functionName);
 };
 
 /**
@@ -271,6 +314,7 @@ const getXRayTraceData = async (
   const traces = await retriableGetTraceDetails({
     traceIds,
     expectedSegmentsCount,
+    functionName: resourceName,
   });
 
   if (!traces) {
@@ -286,9 +330,15 @@ const getXRayTraceData = async (
  * @param options - The options to get the X-Ray trace data, including the start time, resource name, expected traces count, and expected segments count
  */
 const getTraces = async (
-  options: GetXRayTraceIdsOptions & Omit<GetXRayTraceDetailsOptions, 'traceIds'>
+  options: GetXRayTraceIdsOptions &
+    Omit<GetXRayTraceDetailsOptions, 'traceIds' | 'functionName'> & {
+      resourceName: string;
+    }
 ): Promise<EnrichedXRayTraceDocumentParsed[]> => {
-  const traces = await getXRayTraceData(options);
+  const traces = await getXRayTraceData({
+    ...options,
+    functionName: options.resourceName,
+  });
 
   const { resourceName } = options;
 
@@ -305,45 +355,6 @@ const getTraces = async (
   return mainSubsegments;
 };
 
-/**
- * Get the X-Ray trace data for a given resource name without the main subsegments.
- *
- * This is useful when we are testing cases where Active Tracing is disabled and we don't have the main subsegments.
- *
- * @param options - The options to get the X-Ray trace data, including the start time, resource name, expected traces count, and expected segments count
- */
-const getTracesWithoutMainSubsegments = async (
-  options: GetXRayTraceIdsOptions & Omit<GetXRayTraceDetailsOptions, 'traceIds'>
-): Promise<EnrichedXRayTraceDocumentParsed[]> => {
-  const traces = await getXRayTraceData(options);
-
-  const { resourceName } = options;
-
-  const lambdaFunctionSegments = [];
-  for (const trace of traces) {
-    const functionSegment = trace.Segments.find(
-      (segment) => segment.Document.origin === 'AWS::Lambda::Function'
-    );
-
-    if (!functionSegment) {
-      throw new Error(
-        `AWS::Lambda::Function segment not found for ${resourceName}`
-      );
-    }
-
-    const lambdaFunctionSegment = functionSegment.Document;
-    const enrichedSubsegment = {
-      ...lambdaFunctionSegment,
-      subsegments: parseSubsegmentsByName(
-        lambdaFunctionSegment.subsegments ?? []
-      ),
-    };
-    lambdaFunctionSegments.push(enrichedSubsegment);
-  }
-
-  return lambdaFunctionSegments;
-};
-
 export {
   getTraceIds,
   retriableGetTraceIds,
@@ -352,5 +363,4 @@ export {
   findPowertoolsFunctionSegment,
   getTraces,
   parseSubsegmentsByName,
-  getTracesWithoutMainSubsegments,
 };

packages/tracer/package.json

@@ -29,8 +29,7 @@
   "devDependencies": {
     "@aws-lambda-powertools/testing-utils": "file:../testing",
     "@aws-sdk/client-dynamodb": "^3.637.0",
-    "@aws-sdk/client-xray": "^3.637.0",
-    "aws-sdk": "^2.1688.0"
+    "@aws-sdk/client-xray": "^3.637.0"
   },
   "peerDependencies": {
     "@middy/core": "4.x || 5.x"