test(tracer): improve xray trace retrieval

Author: dreamorosi

package-lock.json

@@ -53,6 +53,10 @@
     "./types": {
       "import": "./lib/esm/types.js",
       "require": "./lib/cjs/types.js"
+    },
+    "./utils/xray-traces": {
+      "import": "./lib/esm/xray-traces-utils.js",
+      "require": "./lib/cjs/xray-traces-utils.js"
     }
   },
   "typesVersions": {
@@ -72,6 +76,10 @@
       "context": [
         "lib/cjs/context.d.ts",
         "lib/esm/context.d.ts"
+      ],
+      "utils/xray-traces": [
+        "lib/cjs/xray-traces-utils.d.ts",
+        "lib/esm/xray-traces-utils.d.ts"
       ]
     }
   },
@@ -94,6 +102,10 @@
     "@aws-sdk/client-lambda": "^3.637.0",
     "@smithy/util-utf8": "^3.0.0",
     "aws-cdk-lib": "^2.155.0",
-    "esbuild": "^0.23.1"
+    "esbuild": "^0.23.1",
+    "promise-retry": "^2.0.1"
+  },
+  "devDependencies": {
+    "@types/promise-retry": "^1.1.6"
   }
-}
+}

packages/testing/package.json

@@ -87,6 +87,94 @@ interface TestStackProps {
   stack?: Stack;
 }
 
+// #region X-Ray Trace Utils
+
+type GetXRayTraceIdsOptions = {
+  startTime: Date;
+  resourceName: string;
+  expectedTracesCount: number;
+};
+
+type XRayTraceDocumentParsed = {
+  name: string;
+  id: string;
+  start_time: number;
+  end_time?: number;
+  // This flag may be set if the segment hasn't been fully processed
+  // The trace may have already appeared in the `getTraceSummaries` response
+  // but a segment may still be in_progress
+  in_progress?: boolean;
+  aws?: {
+    request_id: string;
+  };
+  http?: {
+    response: {
+      status: number;
+    };
+  };
+  origin?: string;
+  resource_arn?: string;
+  trace_id?: string;
+  subsegments?: XRayTraceDocumentParsed[];
+  annotations?: {
+    [key: string]: string | boolean | number;
+  };
+  metadata?: {
+    [key: string]: {
+      [key: string]: unknown;
+    };
+  };
+  fault?: boolean;
+  cause?: {
+    working_directory: string;
+    exceptions: {
+      message: string;
+      type: string;
+      remote: boolean;
+      stack: {
+        path: string;
+        line: number;
+        label: string;
+      }[];
+    }[];
+  };
+  exception: {
+    message: string;
+  };
+  error?: boolean;
+};
+
+type XRaySegmentParsed = {
+  Id: string;
+  Document: XRayTraceDocumentParsed;
+};
+
+type XRayTraceParsed = {
+  Id: string;
+  Segments: XRaySegmentParsed[];
+};
+
+type GetXRayTraceDetailsOptions = {
+  /**
+   * The trace IDs to get details for
+   */
+  traceIds: string[];
+  /**
+   * The expected number of segments in each trace
+   */
+  expectedSegmentsCount: number;
+};
+
+/**
+ * Enriched X-Ray trace document parsed with subsegments as a map
+ */
+type EnrichedXRayTraceDocumentParsed = Omit<
+  XRayTraceDocumentParsed,
+  'subsegments'
+> & {
+  subsegments: Map<string, XRayTraceDocumentParsed>;
+};
+
 export type {
   ExtraTestProps,
   TestDynamodbTableProps,
@@ -96,4 +184,10 @@ export type {
   FunctionLog,
   StackNameProps,
   TestStackProps,
+  GetXRayTraceIdsOptions,
+  GetXRayTraceDetailsOptions,
+  XRayTraceDocumentParsed,
+  XRaySegmentParsed,
+  XRayTraceParsed,
+  EnrichedXRayTraceDocumentParsed,
 };