datadog-lambda-js/src/trace/context/extractors/step-function.spec.ts at 89e67d15e5b12aa750cf70e2976cfe10fde436be · DataDog/datadog-lambda-js · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import { StepFunctionContextService } from "../../step-function-service";
import { StepFunctionEventTraceExtractor } from "./step-function";

describe("StepFunctionEventTraceExtractor", () => {
  beforeEach(() => {
    StepFunctionContextService["_instance"] = undefined as any;
  });
  describe("extract", () => {
    const payload = {
      Execution: {
        Id: "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
        Name: "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
        RoleArn:
          "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
        StartTime: "2024-12-04T19:38:04.069Z",
      },
      State: {
        Name: "Lambda Invoke",
        EnteredTime: "2024-12-04T19:38:04.118Z",
        RetryCount: 0,
      },
      StateMachine: {
        Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
        Name: "abhinav-activity-state-machine",
      },
    };

    const redrivePayload = {
      Execution: {
        Id: "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
        Name: "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
        RoleArn:
          "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
        StartTime: "2024-12-04T19:38:04.069Z",
        RedriveCount: 1,
      },
      State: {
        Name: "Lambda Invoke",
        EnteredTime: "2024-12-04T19:38:04.118Z",
        RetryCount: 0,
      },
      StateMachine: {
        Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
        Name: "abhinav-activity-state-machine",
      },
    };
    it("extracts trace context with valid payload", () => {
      // Mimick TraceContextService.extract initialization
      StepFunctionContextService.instance(payload);

      const extractor = new StepFunctionEventTraceExtractor();

      // Payload is sent again for safety in case the instance wasn't previously initialized
      const traceContext = extractor.extract(payload);
      expect(traceContext).not.toBeNull();

      expect(traceContext?.toTraceId()).toBe("435175499815315247");
      expect(traceContext?.toSpanId()).toBe("3929055471293792800");
      expect(traceContext?.sampleMode()).toBe("1");
      expect(traceContext?.source).toBe("event");
    });

    // https://github.com/DataDog/logs-backend/blob/c17618cb552fc369ca40282bae0a65803f82f694/domains/serverless/apps/logs-to-traces-reducer/src/test/resources/test-json-files/stepfunctions/RedriveTest/snapshots/RedriveLambdaSuccessTraceMerging.json#L46
    it("extracts trace context with valid redriven payload", () => {
      // Mimick TraceContextService.extract initialization
      StepFunctionContextService.instance(redrivePayload);

      const extractor = new StepFunctionEventTraceExtractor();

      // Payload is sent again for safety in case the instance wasn't previously initialized
      const traceContext = extractor.extract(redrivePayload);
      expect(traceContext).not.toBeNull();

      expect(traceContext?.toTraceId()).toBe("435175499815315247");
      expect(traceContext?.toSpanId()).toBe("8782364156266188026");
      expect(traceContext?.sampleMode()).toBe("1");
      expect(traceContext?.source).toBe("event");
    });

    it("extracts trace context with valid payload when instance wasn't initialized", () => {
      const extractor = new StepFunctionEventTraceExtractor();

      // This should initialize the current context
      const traceContext = extractor.extract(payload);
      expect(traceContext).not.toBeNull();

      expect(traceContext?.toTraceId()).toBe("435175499815315247");
      expect(traceContext?.toSpanId()).toBe("3929055471293792800");
      expect(traceContext?.sampleMode()).toBe("1");
      expect(traceContext?.source).toBe("event");
    });

    it("extracts trace context with valid legacy lambda payload", () => {
      // Mimick TraceContextService.extract initialization
      StepFunctionContextService.instance({ Payload: payload });

      const extractor = new StepFunctionEventTraceExtractor();

      // Payload is sent again for safety in case the instance wasn't previously initialized
      const traceContext = extractor.extract({ Payload: payload });
      expect(traceContext).not.toBeNull();

      expect(traceContext?.toTraceId()).toBe("435175499815315247");
      expect(traceContext?.toSpanId()).toBe("3929055471293792800");
      expect(traceContext?.sampleMode()).toBe("1");
      expect(traceContext?.source).toBe("event");
    });

    it("returns null when StepFunctionContextService.context is undefined", async () => {
      const extractor = new StepFunctionEventTraceExtractor();

      const traceContext = extractor.extract({});
      expect(traceContext).toBeNull();
    });
  });
});