From 9c863a38f151e84a17e9c44b468fd741617a3a23 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?jordan=20gonz=C3=A1lez?=
 <30836115+duncanista@users.noreply.github.com>
Date: Wed, 7 Feb 2024 13:01:25 -0500
Subject: [PATCH] fix: reading sampling priority will error when not available
 (#483)

* auto keep when sampling priority is not available

* add unit tests

* format
---
 src/trace/span-context-wrapper.spec.ts | 34 +++++++++++++++++++++++++-
 src/trace/span-context-wrapper.ts      |  4 +--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/trace/span-context-wrapper.spec.ts b/src/trace/span-context-wrapper.spec.ts
index 8b4b7693..102d4eb7 100644
--- 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
index 985282a7..b7ec6a4b 100644
--- 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 => {