Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 21 additions & 39 deletions src/trace/cold-start-tracer.spec.ts
Original file line number Diff line number Diff line change
@@ -1,35 +1,17 @@
import { RequireNode } from "../runtime/require-tracer";
import { ColdStartTracerConfig, ColdStartTracer } from "./cold-start-tracer";
import { TracerWrapper, SpanOptions } from "./tracer-wrapper";
import { TracerWrapper } from "./tracer-wrapper";
import { SpanWrapper } from "./span-wrapper";

let mockStartSpan: jest.Mock<any, any>;
let mockFinishSpan: jest.Mock<any, any>;

jest.mock("./tracer-wrapper", () => {
mockFinishSpan = jest.fn();
mockStartSpan = jest.fn().mockImplementation((spanName, spanOptions) => {
describe("ColdStartTracer", () => {
jest.spyOn(TracerWrapper.prototype, "isTracerAvailable", "get").mockReturnValue(true);
const mockFinishSpan = jest.fn();
const startSpanSpy = jest.spyOn(TracerWrapper.prototype, "startSpan").mockImplementation((spanName, spanOptions) => {
return { spanName, spanOptions, finish: mockFinishSpan };
});
class MockTraceWrapper {
get isTraceAvailable() {
return true;
}

constructor() {}

startSpan(spanName: string, spanOptions: SpanOptions): any {
return mockStartSpan(spanName, spanOptions);
}
}
return {
TracerWrapper: MockTraceWrapper,
};
});

describe("ColdStartTracer", () => {
beforeEach(() => {
mockStartSpan.mockClear();
startSpanSpy.mockClear();
mockFinishSpan.mockClear();
});

Expand Down Expand Up @@ -76,17 +58,17 @@ describe("ColdStartTracer", () => {
};
const coldStartTracer = new ColdStartTracer(coldStartConfig);
coldStartTracer.trace(requireNodes);
expect(mockStartSpan).toHaveBeenCalledTimes(5);
expect(startSpanSpy).toHaveBeenCalledTimes(5);
expect(mockFinishSpan).toHaveBeenCalledTimes(5);
const span1 = mockStartSpan.mock.calls[0];
const span1 = startSpanSpy.mock.calls[0];
expect(span1[0]).toEqual("aws.lambda.load");
expect(span1[1].tags).toEqual({
operation_name: "aws.lambda.require",
"resource.name": "my-function-name",
resource_names: "my-function-name",
service: "aws.lambda",
});
const span2 = mockStartSpan.mock.calls[1];
const span2 = startSpanSpy.mock.calls[1];
expect(span2[0]).toEqual("aws.lambda.require");
expect(span2[1].tags).toEqual({
operation_name: "aws.lambda.require",
Expand All @@ -95,7 +77,7 @@ describe("ColdStartTracer", () => {
service: "aws.lambda",
filename: "/var/task/handler.js",
});
const span3 = mockStartSpan.mock.calls[2];
const span3 = startSpanSpy.mock.calls[2];
expect(span3[0]).toEqual("aws.lambda.require_layer");
expect(span3[1].tags).toEqual({
filename: "/opt/nodejs/node_modules/my-child-module.js",
Expand All @@ -104,7 +86,7 @@ describe("ColdStartTracer", () => {
resource_names: "myChildModule",
service: "aws.lambda",
});
const span4 = mockStartSpan.mock.calls[3];
const span4 = startSpanSpy.mock.calls[3];
expect(span4[0]).toEqual("aws.lambda.require_core_module");
expect(span4[1].tags).toEqual({
filename: "http",
Expand All @@ -113,7 +95,7 @@ describe("ColdStartTracer", () => {
resource_names: "myCoreModule",
service: "aws.lambda",
});
const span5 = mockStartSpan.mock.calls[4];
const span5 = startSpanSpy.mock.calls[4];
expect(span5[0]).toEqual("aws.lambda.require_runtime");
expect(span5[1].tags).toEqual({
filename: "/var/runtime/aws-sdk",
Expand Down Expand Up @@ -167,17 +149,17 @@ describe("ColdStartTracer", () => {
};
const coldStartTracer = new ColdStartTracer(coldStartConfig);
coldStartTracer.trace(requireNodes);
expect(mockStartSpan).toHaveBeenCalledTimes(3);
expect(startSpanSpy).toHaveBeenCalledTimes(3);
expect(mockFinishSpan).toHaveBeenCalledTimes(3);
const span1 = mockStartSpan.mock.calls[0];
const span1 = startSpanSpy.mock.calls[0];
expect(span1[0]).toEqual("aws.lambda.load");
expect(span1[1].tags).toEqual({
operation_name: "aws.lambda.require",
"resource.name": "my-function-name",
resource_names: "my-function-name",
service: "aws.lambda",
});
const span2 = mockStartSpan.mock.calls[1];
const span2 = startSpanSpy.mock.calls[1];
expect(span2[0]).toEqual("aws.lambda.require");
expect(span2[1].tags).toEqual({
operation_name: "aws.lambda.require",
Expand All @@ -186,7 +168,7 @@ describe("ColdStartTracer", () => {
service: "aws.lambda",
filename: "/var/task/handler.js",
});
const span3 = mockStartSpan.mock.calls[2];
const span3 = startSpanSpy.mock.calls[2];
expect(span3[0]).toEqual("aws.lambda.require_runtime");
expect(span3[1].tags).toEqual({
filename: "/var/runtime/aws-sdk",
Expand Down Expand Up @@ -239,9 +221,9 @@ describe("ColdStartTracer", () => {
};
const coldStartTracer = new ColdStartTracer(coldStartConfig);
coldStartTracer.trace(requireNodes);
expect(mockStartSpan).toHaveBeenCalledTimes(4);
expect(startSpanSpy).toHaveBeenCalledTimes(4);
expect(mockFinishSpan).toHaveBeenCalledTimes(4);
const span1 = mockStartSpan.mock.calls[0];
const span1 = startSpanSpy.mock.calls[0];
expect(span1[0]).toEqual("aws.lambda.require");
expect(span1[1].tags).toEqual({
operation_name: "aws.lambda.require",
Expand All @@ -250,7 +232,7 @@ describe("ColdStartTracer", () => {
service: "aws.lambda",
filename: "/var/task/handler.js",
});
const span2 = mockStartSpan.mock.calls[1];
const span2 = startSpanSpy.mock.calls[1];
expect(span2[0]).toEqual("aws.lambda.require_layer");
expect(span2[1].tags).toEqual({
filename: "/opt/nodejs/node_modules/my-child-module.js",
Expand All @@ -259,7 +241,7 @@ describe("ColdStartTracer", () => {
resource_names: "myChildModule",
service: "aws.lambda",
});
const span3 = mockStartSpan.mock.calls[2];
const span3 = startSpanSpy.mock.calls[2];
expect(span3[0]).toEqual("aws.lambda.require_core_module");
expect(span3[1].tags).toEqual({
filename: "http",
Expand All @@ -268,7 +250,7 @@ describe("ColdStartTracer", () => {
resource_names: "myCoreModule",
service: "aws.lambda",
});
const span4 = mockStartSpan.mock.calls[3];
const span4 = startSpanSpy.mock.calls[3];
expect(span4[0]).toEqual("aws.lambda.require_runtime");
expect(span4[1].tags).toEqual({
filename: "/var/runtime/aws-sdk",
Expand Down
68 changes: 34 additions & 34 deletions src/trace/context/extractor.spec.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import crypto from "node:crypto";
import { Context, EventBridgeEvent, KinesisStreamEvent, SNSEvent, SQSEvent } from "aws-lambda";
import { TraceConfig } from "../listener";
import { TracerWrapper } from "../tracer-wrapper";
Expand All @@ -21,19 +22,23 @@ import {
import { StepFunctionContextService } from "../step-function-service";
import { SpanContextWrapper } from "../span-context-wrapper";

let mockSpanContext: any = null;
let sentSegment: any;
let closedSocket = false;
const mockController: {
sentSegment?: any;
closedSocket: boolean;
mockSpanContext: any;
} = {
closedSocket: false,
mockSpanContext: null,
};

// Mocking extract is needed, due to dd-trace being a No-op
// if the detected environment is testing. This is expected, since
// we don't want to test dd-trace extraction, but our components.
const ddTrace = require("dd-trace");
jest.mock("dd-trace", () => {
return {
...ddTrace,
...jest.requireActual("dd-trace"),
_tracer: { _service: {} },
extract: (_carrier: any, _headers: any) => mockSpanContext,
extract: (_carrier: any, _headers: any) => mockController.mockSpanContext,
};
});

Expand All @@ -48,28 +53,23 @@ jest.mock("dgram", () => ({
address: string,
callback: (error: string | undefined, bytes: number) => void,
) => {
sentSegment = message;
mockController.sentSegment = message;
callback(undefined, 1);
},
close: () => {
closedSocket = true;
mockController.closedSocket = true;
},
};
},
}));

jest.mock("crypto", () => {
return {
randomBytes: () => "11111",
};
});

const spyTracerWrapper = jest.spyOn(TracerWrapper.prototype, "extract");

describe("TraceContextExtractor", () => {
describe("extract", () => {
beforeEach(() => {
mockSpanContext = null;
jest.spyOn(crypto, "randomBytes").mockImplementation(() => "11111");
mockController.mockSpanContext = null;
StepFunctionContextService["_instance"] = undefined as any;
});

Expand Down Expand Up @@ -153,7 +153,7 @@ describe("TraceContextExtractor", () => {
});
// HTTP event
it("extracts trace context from HTTP headers", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "4110911582297405551",
toSpanId: () => "797643193680388251",
_sampling: {
Expand Down Expand Up @@ -188,7 +188,7 @@ describe("TraceContextExtractor", () => {

// SNS message event (String Value)
it("extracts trace context from SNS event - String Value", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "6966585609680374559",
toSpanId: () => "4297634551783724228",
_sampling: {
Expand Down Expand Up @@ -250,7 +250,7 @@ describe("TraceContextExtractor", () => {

// SNS message event (Binary Value)
it("extracts trace context from SNS event - Binary Value", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "7102291628443134919",
toSpanId: () => "4247550101648618618",
_sampling: {
Expand Down Expand Up @@ -311,7 +311,7 @@ describe("TraceContextExtractor", () => {

// SNS message delivered to SQS queue event (String Value)
it("extracts trace context from SNS to SQS event - String Value", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "2776434475358637757",
toSpanId: () => "4493917105238181843",
_sampling: {
Expand Down Expand Up @@ -362,7 +362,7 @@ describe("TraceContextExtractor", () => {

// SNS message delivered to SQS queue event (Binary Value)
it("extracts trace context from SNS to SQS event - Binary Value", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "7102291628443134919",
toSpanId: () => "4247550101648618618",
_sampling: {
Expand Down Expand Up @@ -412,7 +412,7 @@ describe("TraceContextExtractor", () => {

// EventBridge message delivered to SQS queue event
it("extracts trace context from EventBridge to SQS event", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "7379586022458917877",
toSpanId: () => "2644033662113726488",
_sampling: {
Expand Down Expand Up @@ -464,7 +464,7 @@ describe("TraceContextExtractor", () => {

// AppSync event
it("extracts trace context from AppSync event", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "797643193680388254",
toSpanId: () => "4110911582297405557",
_sampling: {
Expand Down Expand Up @@ -504,7 +504,7 @@ describe("TraceContextExtractor", () => {

// SQS queue message event
it("extracts trace context from SQS event", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "4555236104497098341",
toSpanId: () => "3369753143434738315",
_sampling: {
Expand Down Expand Up @@ -562,7 +562,7 @@ describe("TraceContextExtractor", () => {

// Kinesis stream event
it("extracts trace context from Kinesis event", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "667309514221035538",
toSpanId: () => "1350735035497811828",
_sampling: {
Expand Down Expand Up @@ -612,7 +612,7 @@ describe("TraceContextExtractor", () => {

// EventBridge message event
it("extracts trace context from EventBridge event", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "5827606813695714842",
toSpanId: () => "4726693487091824375",
_sampling: {
Expand Down Expand Up @@ -745,7 +745,7 @@ describe("TraceContextExtractor", () => {

describe("lambda context", () => {
it("extracts trace context from Lambda context", async () => {
mockSpanContext = {
mockController.mockSpanContext = {
toTraceId: () => "667309514221035538",
toSpanId: () => "1350735035497811828",
_sampling: {
Expand Down Expand Up @@ -926,8 +926,8 @@ describe("getTraceEventExtractor", () => {
describe("addTraceContextToXray", () => {
beforeEach(() => {
StepFunctionContextService["_instance"] = undefined as any;
sentSegment = undefined;
closedSocket = false;
mockController.sentSegment = undefined;
mockController.closedSocket = false;
process.env["_X_AMZN_TRACE_ID"] = undefined;
process.env["AWS_XRAY_DAEMON_ADDRESS"] = undefined;
});
Expand Down Expand Up @@ -979,10 +979,10 @@ describe("getTraceEventExtractor", () => {

extractor["addTraceContextToXray"](spanContext);

expect(sentSegment).toBeInstanceOf(Buffer);
expect(closedSocket).toBeTruthy();
expect(mockController.sentSegment).toBeInstanceOf(Buffer);
expect(mockController.closedSocket).toBeTruthy();

const sentMessage = sentSegment.toString();
const sentMessage = mockController.sentSegment.toString();
expect(sentMessage).toEqual(
'{"format": "json", "version": 1}\n{"id":"11111","trace_id":"1-5e272390-8c398be037738dc042009320","parent_id":"94ae789b969f1cc5","name":"datadog-metadata","start_time":1487076708,"end_time":1487076708,"type":"subsegment","metadata":{"datadog":{"root_span_metadata":{"execution_id":"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf","redrive_count":"0","retry_count":"2","state_entered_time":"2022-12-08T21:08:19.224Z","state_name":"step-one"}}}}',
);
Expand Down Expand Up @@ -1010,10 +1010,10 @@ describe("getTraceEventExtractor", () => {

extractor["addTraceContextToXray"](spanContext);

expect(sentSegment).toBeInstanceOf(Buffer);
expect(closedSocket).toBeTruthy();
expect(mockController.sentSegment).toBeInstanceOf(Buffer);
expect(mockController.closedSocket).toBeTruthy();

const sentMessage = sentSegment.toString();
const sentMessage = mockController.sentSegment.toString();
expect(sentMessage).toEqual(
'{"format": "json", "version": 1}\n{"id":"11111","trace_id":"1-5e272390-8c398be037738dc042009320","parent_id":"94ae789b969f1cc5","name":"datadog-metadata","start_time":1487076708,"end_time":1487076708,"type":"subsegment","metadata":{"datadog":{"trace":{"trace-id":"4110911582297405551","parent-id":"797643193680388251","sampling-priority":"2"}}}}',
);
Expand Down
3 changes: 1 addition & 2 deletions src/trace/context/extractors/app-sync.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,9 @@ let mockSpanContext: any = null;
// Mocking extract is needed, due to dd-trace being a No-op
// if the detected environment is testing. This is expected, since
// we don't want to test dd-trace extraction, but our components.
const ddTrace = require("dd-trace");
jest.mock("dd-trace", () => {
return {
...ddTrace,
...jest.requireActual("dd-trace"),
_tracer: { _service: {} },
extract: (_carrier: any, _headers: any) => mockSpanContext,
};
Expand Down
3 changes: 1 addition & 2 deletions src/trace/context/extractors/event-bridge-sqs.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,9 @@ let mockSpanContext: any = null;
// Mocking extract is needed, due to dd-trace being a No-op
// if the detected environment is testing. This is expected, since
// we don't want to test dd-trace extraction, but our components.
const ddTrace = require("dd-trace");
jest.mock("dd-trace", () => {
return {
...ddTrace,
...jest.requireActual("dd-trace"),
_tracer: { _service: {} },
extract: (_carrier: any, _headers: any) => mockSpanContext,
};
Expand Down
Loading