diff --git a/src/trace/cold-start-tracer.spec.ts b/src/trace/cold-start-tracer.spec.ts index ebf834b6..78827d86 100644 --- a/src/trace/cold-start-tracer.spec.ts +++ b/src/trace/cold-start-tracer.spec.ts @@ -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; -let mockFinishSpan: jest.Mock; - -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(); }); @@ -76,9 +58,9 @@ 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", @@ -86,7 +68,7 @@ describe("ColdStartTracer", () => { 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", @@ -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", @@ -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", @@ -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", @@ -167,9 +149,9 @@ 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", @@ -177,7 +159,7 @@ describe("ColdStartTracer", () => { 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", @@ -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", @@ -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", @@ -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", @@ -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", @@ -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", diff --git a/src/trace/context/extractor.spec.ts b/src/trace/context/extractor.spec.ts index 37cae4a3..f41a645e 100644 --- a/src/trace/context/extractor.spec.ts +++ b/src/trace/context/extractor.spec.ts @@ -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"; @@ -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, }; }); @@ -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; }); @@ -153,7 +153,7 @@ describe("TraceContextExtractor", () => { }); // HTTP event it("extracts trace context from HTTP headers", async () => { - mockSpanContext = { + mockController.mockSpanContext = { toTraceId: () => "4110911582297405551", toSpanId: () => "797643193680388251", _sampling: { @@ -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: { @@ -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: { @@ -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: { @@ -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: { @@ -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: { @@ -464,7 +464,7 @@ describe("TraceContextExtractor", () => { // AppSync event it("extracts trace context from AppSync event", async () => { - mockSpanContext = { + mockController.mockSpanContext = { toTraceId: () => "797643193680388254", toSpanId: () => "4110911582297405557", _sampling: { @@ -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: { @@ -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: { @@ -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: { @@ -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: { @@ -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; }); @@ -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"}}}}', ); @@ -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"}}}}', ); diff --git a/src/trace/context/extractors/app-sync.spec.ts b/src/trace/context/extractors/app-sync.spec.ts index a1c217d9..f15611f8 100644 --- a/src/trace/context/extractors/app-sync.spec.ts +++ b/src/trace/context/extractors/app-sync.spec.ts @@ -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, }; diff --git a/src/trace/context/extractors/event-bridge-sqs.spec.ts b/src/trace/context/extractors/event-bridge-sqs.spec.ts index 24ca16fa..a457e479 100644 --- a/src/trace/context/extractors/event-bridge-sqs.spec.ts +++ b/src/trace/context/extractors/event-bridge-sqs.spec.ts @@ -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, }; diff --git a/src/trace/context/extractors/event-bridge.spec.ts b/src/trace/context/extractors/event-bridge.spec.ts index 346b89c4..f3d0a3e5 100644 --- a/src/trace/context/extractors/event-bridge.spec.ts +++ b/src/trace/context/extractors/event-bridge.spec.ts @@ -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, }; diff --git a/src/trace/context/extractors/http.spec.ts b/src/trace/context/extractors/http.spec.ts index 1cb7b3f4..01b42f67 100644 --- a/src/trace/context/extractors/http.spec.ts +++ b/src/trace/context/extractors/http.spec.ts @@ -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, }; diff --git a/src/trace/context/extractors/kinesis.spec.ts b/src/trace/context/extractors/kinesis.spec.ts index 3f340359..87e65c88 100644 --- a/src/trace/context/extractors/kinesis.spec.ts +++ b/src/trace/context/extractors/kinesis.spec.ts @@ -15,10 +15,9 @@ jest.mock("dd-trace/packages/dd-trace/src/datastreams/checkpointer", () => { // 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, dataStreamsCheckpointer: mockDataStreamsCheckpointer, diff --git a/src/trace/context/extractors/lambda-context.spec.ts b/src/trace/context/extractors/lambda-context.spec.ts index fb3ef4f4..e935c4c9 100644 --- a/src/trace/context/extractors/lambda-context.spec.ts +++ b/src/trace/context/extractors/lambda-context.spec.ts @@ -9,7 +9,7 @@ let mockSpanContext: any = null; const ddTrace = require("dd-trace"); jest.mock("dd-trace", () => { return { - ...ddTrace, + ...jest.requireActual("dd-trace"), _tracer: { _service: {} }, extract: (_carrier: any, _headers: any) => mockSpanContext, }; diff --git a/src/trace/context/extractors/sns-sqs.spec.ts b/src/trace/context/extractors/sns-sqs.spec.ts index 92f77049..93f51948 100644 --- a/src/trace/context/extractors/sns-sqs.spec.ts +++ b/src/trace/context/extractors/sns-sqs.spec.ts @@ -16,10 +16,9 @@ jest.mock("dd-trace/packages/dd-trace/src/datastreams/checkpointer", () => { // 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, dataStreamsCheckpointer: mockDataStreamsCheckpointer, diff --git a/src/trace/context/extractors/sns.spec.ts b/src/trace/context/extractors/sns.spec.ts index 8932ee08..3ac18a99 100644 --- a/src/trace/context/extractors/sns.spec.ts +++ b/src/trace/context/extractors/sns.spec.ts @@ -17,10 +17,9 @@ jest.mock("dd-trace/packages/dd-trace/src/datastreams/checkpointer", () => { // 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, dataStreamsCheckpointer: mockDataStreamsCheckpointer, diff --git a/src/trace/context/extractors/sqs.spec.ts b/src/trace/context/extractors/sqs.spec.ts index ef351ad2..bbf9cdd8 100644 --- a/src/trace/context/extractors/sqs.spec.ts +++ b/src/trace/context/extractors/sqs.spec.ts @@ -59,10 +59,9 @@ const makeRecord = ({ // 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, dataStreamsCheckpointer: mockDataStreamsCheckpointer, diff --git a/src/trace/listener.spec.ts b/src/trace/listener.spec.ts index 1e8dc189..8b354e17 100644 --- a/src/trace/listener.spec.ts +++ b/src/trace/listener.spec.ts @@ -10,63 +10,25 @@ import { DATADOG_SAMPLING_PRIORITY_HEADER, DATADOG_TRACE_ID_HEADER, } from "./context/extractor"; +import { TracerWrapper } from "./tracer-wrapper"; -let mockWrap: jest.Mock; -let mockExtract: jest.Mock; -let mockSpanContextWrapper: any; -let mockSpanContext: any; -let mockTraceSource: TraceSource | undefined = undefined; - -jest.mock("./tracer-wrapper", () => { - mockWrap = jest.fn().mockImplementation((name, options, func) => func); - mockExtract = jest.fn().mockImplementation((val) => val); - class MockTraceWrapper { - public isTracerAvailable = true; - - constructor() {} - - wrap(name: any, options: any, fn: any): any { - return mockWrap(name, options, fn); - } - - extract(event: any): any { - return mockExtract(event); - } - - startSpan(name: any, options: any): any { - return { - toSpanId: () => "mockSpanId", - toTraceId: () => "mockTraceId", - finish: jest.fn(), - setTag: jest.fn(), - }; - } - - injectSpan(span: any): any { - return { - [DATADOG_PARENT_ID_HEADER]: span.toSpanId(), - [DATADOG_TRACE_ID_HEADER]: span.toTraceId(), - [DATADOG_SAMPLING_PRIORITY_HEADER]: 1, - [parentSpanFinishTimeHeader]: 1661189936981, - }; - } - } - return { - TracerWrapper: MockTraceWrapper, - }; -}); +const mockController: { + mockSpanContext?: any; + mockSpanContextWrapper?: any; + mockTraceSource?: TraceSource; +} = {}; jest.mock("./trace-context-service", () => { class MockTraceContextService { extract(event: any, context: Context): SpanContextWrapper { - return mockSpanContextWrapper; + return mockController.mockSpanContextWrapper; } get traceSource() { - return mockTraceSource; + return mockController.mockTraceSource; } get currentTraceContext() { - return mockSpanContextWrapper; + return mockController.mockSpanContextWrapper; } reset() { // mocking @@ -79,6 +41,22 @@ jest.mock("./trace-context-service", () => { }); describe("TraceListener", () => { + jest.spyOn(TracerWrapper.prototype, "isTracerAvailable", "get").mockReturnValue(true); + jest.spyOn(TracerWrapper.prototype, "extract").mockImplementation((val) => val); + jest.spyOn(TracerWrapper.prototype, "startSpan").mockReturnValue({ + toSpanId: () => "mockSpanId", + toTraceId: () => "mockTraceId", + finish: jest.fn(), + setTag: jest.fn(), + }); + jest.spyOn(TracerWrapper.prototype, "injectSpan").mockImplementation((span) => ({ + [DATADOG_PARENT_ID_HEADER]: span.toSpanId(), + [DATADOG_TRACE_ID_HEADER]: span.toTraceId(), + [DATADOG_SAMPLING_PRIORITY_HEADER]: 1, + [parentSpanFinishTimeHeader]: 1661189936981, + })); + const wrapSpy = jest.spyOn(TracerWrapper.prototype, "wrap").mockImplementation((name, options, func) => func); + let oldEnv: any; const defaultConfig = { autoPatchHTTP: true, @@ -108,11 +86,10 @@ describe("TraceListener", () => { invokedFunctionArn: "arn:aws:lambda:us-east-1:123456789101:function:my-lambda:1", }; beforeEach(() => { - mockWrap.mockClear(); - mockExtract.mockClear(); - mockSpanContext = undefined; - mockSpanContextWrapper = undefined; - mockTraceSource = undefined; + wrapSpy.mockClear(); + mockController.mockSpanContext = undefined; + mockController.mockSpanContextWrapper = undefined; + mockController.mockTraceSource = undefined; oldEnv = process.env; process.env = { ...oldEnv }; delete process.env.DD_SERVICE; @@ -130,7 +107,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -153,16 +130,16 @@ describe("TraceListener", () => { it("wraps dd-trace span around invocation, with trace context from event", async () => { const listener = new TraceListener(defaultConfig); - mockTraceSource = TraceSource.Event; - mockSpanContext = { + mockController.mockTraceSource = TraceSource.Event; + mockController.mockSpanContext = { toTraceId: () => "4110911582297405551", toSpanId: () => "797643193680388251", _sampling: { priority: "2", }, }; - mockSpanContextWrapper = { - spanContext: mockSpanContext, + mockController.mockSpanContextWrapper = { + spanContext: mockController.mockSpanContext, }; await listener.onStartInvocation({}, context as any); const unwrappedFunc = () => {}; @@ -170,7 +147,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -187,7 +164,7 @@ describe("TraceListener", () => { dd_trace: ddtraceVersion, }, type: "serverless", - childOf: mockSpanContext, + childOf: mockController.mockSpanContext, }, unwrappedFunc, ); @@ -195,7 +172,7 @@ describe("TraceListener", () => { it("wraps dd-trace span around invocation, without trace context from xray", async () => { const listener = new TraceListener(defaultConfig); - mockTraceSource = TraceSource.Xray; + mockController.mockTraceSource = TraceSource.Xray; await listener.onStartInvocation({}, context as any); const unwrappedFunc = () => {}; @@ -203,7 +180,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -226,16 +203,16 @@ describe("TraceListener", () => { it("wraps dd-trace span around invocation, with trace context from xray when mergeDatadogXrayTraces is enabled", async () => { const listener = new TraceListener({ ...defaultConfig, mergeDatadogXrayTraces: true }); - mockTraceSource = TraceSource.Xray; - mockSpanContext = { + mockController.mockTraceSource = TraceSource.Xray; + mockController.mockSpanContext = { toTraceId: () => "4110911582297405551", toSpanId: () => "797643193680388251", _sampling: { priority: "2", }, }; - mockSpanContextWrapper = { - spanContext: mockSpanContext, + mockController.mockSpanContextWrapper = { + spanContext: mockController.mockSpanContext, }; await listener.onStartInvocation({}, context as any); @@ -244,7 +221,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -261,7 +238,7 @@ describe("TraceListener", () => { dd_trace: ddtraceVersion, }, type: "serverless", - childOf: mockSpanContext, + childOf: mockController.mockSpanContext, }, unwrappedFunc, ); @@ -275,7 +252,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -304,7 +281,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -327,18 +304,18 @@ describe("TraceListener", () => { it("wraps dd-trace span around invocation with Step Function context", async () => { const listener = new TraceListener(defaultConfig); - mockTraceSource = TraceSource.Event; + mockController.mockTraceSource = TraceSource.Event; // Mock Step Function context with deterministic trace IDs - mockSpanContext = { + mockController.mockSpanContext = { toTraceId: () => "512d06a10e5e34cb", // Hex converted to decimal would be different toSpanId: () => "7069a031ef9ad2cc", _sampling: { priority: "1", }, }; - mockSpanContextWrapper = { - spanContext: mockSpanContext, + mockController.mockSpanContextWrapper = { + spanContext: mockController.mockSpanContext, }; const stepFunctionSQSEvent = { @@ -378,7 +355,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -409,7 +386,7 @@ describe("TraceListener", () => { it("injects authorizer context if it exists", async () => { const listener = new TraceListener(defaultConfig); - mockTraceSource = TraceSource.Event; + mockController.mockTraceSource = TraceSource.Event; const inferredSpan = new SpanWrapper( { toSpanId: () => { @@ -445,7 +422,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", { resource: "my-Lambda", @@ -474,11 +451,10 @@ describe("TraceListener", () => { }; beforeEach(() => { - mockWrap.mockClear(); - mockExtract.mockClear(); - mockSpanContext = undefined; - mockSpanContextWrapper = undefined; - mockTraceSource = undefined; + wrapSpy.mockClear(); + mockController.mockSpanContext = undefined; + mockController.mockSpanContextWrapper = undefined; + mockController.mockTraceSource = undefined; process.env = { ...oldEnv }; // Restore original environment variables delete process.env.DD_SERVICE; // Ensure DD_SERVICE doesn't interfere delete process.env.DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED; @@ -497,7 +473,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", expect.objectContaining({ service: "aws.lambda", @@ -515,7 +491,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", expect.objectContaining({ service: "aws.lambda", @@ -533,7 +509,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", expect.objectContaining({ service: lambdaContext.functionName, @@ -551,7 +527,7 @@ describe("TraceListener", () => { wrappedFunc(); await listener.onCompleteInvocation(); - expect(mockWrap).toHaveBeenCalledWith( + expect(wrapSpy).toHaveBeenCalledWith( "aws.lambda", expect.objectContaining({ service: lambdaContext.functionName, diff --git a/tsconfig.json b/tsconfig.json index 256e1a0f..7d7320d8 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -2,7 +2,7 @@ "compilerOptions": { /* Basic Options */ // "incremental": true, /* Enable incremental compilation */ - "target": "es5" /* Specify ECMAScript target version: 'ES3' (default), 'ES5', 'ES2015', 'ES2016', 'ES2017', 'ES2018', 'ES2019' or 'ESNEXT'. */, + "target": "es2021" /* Specify ECMAScript target version: 'ES3' (default), 'ES5', 'ES2015', 'ES2016', 'ES2017', 'ES2018', 'ES2019' or 'ESNEXT'. */, "module": "commonjs" /* Specify module code generation: 'none', 'commonjs', 'amd', 'system', 'umd', 'es2015', or 'ESNext'. */, "lib": [ "es2015",