Merge bf0b1ae into a481762

DataDog · Feb 7, 2024 · 226cee8 · 226cee8
2 parents a481762 + bf0b1ae
commit 226cee8
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/src/trace/span-context-wrapper.spec.ts b/src/trace/span-context-wrapper.spec.ts
@@ -1,5 +1,5 @@
 import { SpanContextWrapper } from "./span-context-wrapper";
-import { TraceSource } from "./trace-context-service";
+import { SampleMode, TraceSource } from "./trace-context-service";
 
 describe("SpanContextWrapper", () => {
   beforeEach(() => {});
@@ -20,4 +20,36 @@ describe("SpanContextWrapper", () => {
     expect(spanContext?.spanContext._traceId.toArray()).toEqual([121, 177, 18, 140, 107, 104, 185, 240]);
     expect(spanContext?.spanContext._spanId.toArray()).toEqual([96, 4, 217, 174, 182, 29, 220, 172]);
   });
+
+  describe("sampleMode", () => {
+    it("should return AUTO_KEEP when sampling priority is not available in spanContext", () => {
+      const spanContext = new SpanContextWrapper(
+        {
+          toSpanId: () => "1234",
+          toTraceId: () => "5678",
+          _sampling: {},
+        },
+        TraceSource.Event,
+      );
+
+      const sampleMode = spanContext.sampleMode();
+      expect(sampleMode).toBe(SampleMode.AUTO_KEEP);
+      expect(sampleMode.toString()).toBe("1");
+    });
+
+    it("should return sampling priority when available in spanContext", () => {
+      const spanContext = new SpanContextWrapper(
+        {
+          toSpanId: () => "1234",
+          toTraceId: () => "5678",
+          _sampling: { priority: 2 },
+        },
+        TraceSource.Event,
+      );
+
+      const sampleMode = spanContext.sampleMode();
+      expect(sampleMode).toBe(2);
+      expect(sampleMode.toString()).toBe("2");
+    });
+  });
 });
diff --git a/src/trace/span-context-wrapper.ts b/src/trace/span-context-wrapper.ts
@@ -1,5 +1,5 @@
 import { logDebug } from "../utils";
-import { TraceContext, TraceSource } from "./trace-context-service";
+import { SampleMode, TraceContext, TraceSource } from "./trace-context-service";
 import { SpanContext } from "./tracer-wrapper";
 
 /**
@@ -18,7 +18,7 @@ export class SpanContextWrapper implements SpanContext {
   }
 
   public sampleMode(): number {
-    return this.spanContext._sampling?.priority;
+    return this.spanContext._sampling?.priority ?? SampleMode.AUTO_KEEP;
   }
 
   public toString = (): string => {