Skip to content

Commit c5130af

Browse files
ryanstratryanstrat
committed
Clean up extractor.ts and add tests for extractor-utils.ts
1 parent ef47d0b commit c5130af

File tree

6 files changed

+7347
-5244
lines changed

6 files changed

+7347
-5244
lines changed

‎src/trace/context/extractor-utils.spec.ts‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
import { TracerWrapper } from "../tracer-wrapper";
2+
import { extractTraceContext, extractFromAWSTraceHeader } from "./extractor-utils";
3+
import { StepFunctionContextService } from "../step-function-service";
4+
5+
describe("extractor-utils", () => {
6+
beforeEach(() => {
7+
StepFunctionContextService["_instance"] = undefined as any;
8+
});
9+
describe("extractTraceContext", () => {
10+
it("returns span context when tracer wrapper successfully extracts from headers", () => {
11+
const legacyStepFunctionEvent = {
12+
Execution: {
13+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
14+
Input: {
15+
MyInput: "MyValue",
16+
},
17+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
18+
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
19+
RedriveCount: 0,
20+
StartTime: "2022-12-08T21:08:17.924Z",
21+
},
22+
State: {
23+
Name: "step-one",
24+
EnteredTime: "2022-12-08T21:08:19.224Z",
25+
RetryCount: 2,
26+
},
27+
StateMachine: {
28+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
29+
Name: "my-state-machine",
30+
},
31+
};
32+
33+
const tracerWrapper = new TracerWrapper();
34+
const result = extractTraceContext(legacyStepFunctionEvent, tracerWrapper);
35+
36+
// Should return a span context from Step Function context since headers extraction fails
37+
expect(result).not.toBeNull();
38+
});
39+
40+
it("returns null when no trace context can be extracted", () => {
41+
const emptyEvent = {
42+
someOtherProperty: "value",
43+
};
44+
45+
const tracerWrapper = new TracerWrapper();
46+
const result = extractTraceContext(emptyEvent, tracerWrapper);
47+
48+
expect(result).toBeNull();
49+
});
50+
51+
it("extracts context from LambdaRootStepFunctionContext", () => {
52+
const lambdaRootStepFunctionEvent = {
53+
_datadog: {
54+
Execution: {
55+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
56+
Input: {
57+
MyInput: "MyValue",
58+
},
59+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
60+
RoleArn:
61+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
62+
RedriveCount: 0,
63+
StartTime: "2022-12-08T21:08:17.924Z",
64+
},
65+
State: {
66+
Name: "step-one",
67+
EnteredTime: "2022-12-08T21:08:19.224Z",
68+
RetryCount: 2,
69+
},
70+
StateMachine: {
71+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
72+
Name: "my-state-machine",
73+
},
74+
"x-datadog-trace-id": "10593586103637578129",
75+
"x-datadog-tags": "_dd.p.dm=-0,_dd.p.tid=6734e7c300000000",
76+
"serverless-version": "v1",
77+
},
78+
};
79+
80+
const tracerWrapper = new TracerWrapper();
81+
const result = extractTraceContext(lambdaRootStepFunctionEvent, tracerWrapper);
82+
83+
expect(result).not.toBeNull();
84+
});
85+
86+
it("extracts context from NestedStepFunctionContext", () => {
87+
const nestedStepFunctionEvent = {
88+
_datadog: {
89+
Execution: {
90+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
91+
Input: {
92+
MyInput: "MyValue",
93+
},
94+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
95+
RoleArn:
96+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
97+
RedriveCount: 0,
98+
StartTime: "2022-12-08T21:08:17.924Z",
99+
},
100+
State: {
101+
Name: "step-one",
102+
EnteredTime: "2022-12-08T21:08:19.224Z",
103+
RetryCount: 2,
104+
},
105+
StateMachine: {
106+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
107+
Name: "my-state-machine",
108+
},
109+
RootExecutionId:
110+
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:a1b2c3d4-e5f6-7890-1234-56789abcdef0:9f8e7d6c-5b4a-3c2d-1e0f-123456789abc",
111+
"serverless-version": "v1",
112+
},
113+
};
114+
115+
const tracerWrapper = new TracerWrapper();
116+
const result = extractTraceContext(nestedStepFunctionEvent, tracerWrapper);
117+
118+
expect(result).not.toBeNull();
119+
});
120+
121+
it("extracts context from legacy lambda StepFunctionContext", () => {
122+
const event = {
123+
Payload: {
124+
Execution: {
125+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
126+
Input: {
127+
MyInput: "MyValue",
128+
},
129+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
130+
RoleArn:
131+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
132+
RedriveCount: 0,
133+
StartTime: "2022-12-08T21:08:17.924Z",
134+
},
135+
State: {
136+
Name: "step-one",
137+
EnteredTime: "2022-12-08T21:08:19.224Z",
138+
RetryCount: 2,
139+
},
140+
StateMachine: {
141+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
142+
Name: "my-state-machine",
143+
},
144+
},
145+
};
146+
147+
const tracerWrapper = new TracerWrapper();
148+
const result = extractTraceContext(event, tracerWrapper);
149+
150+
expect(result).not.toBeNull();
151+
});
152+
});
153+
154+
describe("extractFromAWSTraceHeader", () => {
155+
it("returns null when AWS trace header is invalid", () => {
156+
const invalidHeader = "invalid-header";
157+
const eventType = "SQS";
158+
159+
const result = extractFromAWSTraceHeader(invalidHeader, eventType);
160+
161+
expect(result).toBeNull();
162+
});
163+
});
164+
});

0 commit comments

Comments
 (0)