test(sampling): builtin samplers should forward traceState

raphael-theriault-swi · Jan 12, 2023 · e008229 · e008229
1 parent 467b282
commit e008229
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 1 deletion.
diff --git a/packages/opentelemetry-sdk-trace-base/test/common/sampler/AlwaysOffSampler.test.ts b/packages/opentelemetry-sdk-trace-base/test/common/sampler/AlwaysOffSampler.test.ts
@@ -30,4 +30,18 @@ describe('AlwaysOffSampler', () => {
       traceState: undefined,
     });
   });
+
+  it('should forward the traceState', () => {
+    const sampler = new AlwaysOffSampler();
+    const traceState = api.createTraceState();
+    assert.strictEqual(
+      sampler.shouldSample(
+        api.trace.setSpanContext(api.ROOT_CONTEXT, {
+          ...api.INVALID_SPAN_CONTEXT,
+          traceState,
+        })
+      ).traceState,
+      traceState
+    );
+  });
 });
diff --git a/packages/opentelemetry-sdk-trace-base/test/common/sampler/AlwaysOnSampler.test.ts b/packages/opentelemetry-sdk-trace-base/test/common/sampler/AlwaysOnSampler.test.ts
@@ -30,4 +30,18 @@ describe('AlwaysOnSampler', () => {
       traceState: undefined,
     });
   });
+
+  it('should forward the traceState', () => {
+    const sampler = new AlwaysOnSampler();
+    const traceState = api.createTraceState();
+    assert.deepStrictEqual(
+      sampler.shouldSample(
+        api.trace.setSpanContext(api.ROOT_CONTEXT, {
+          ...api.INVALID_SPAN_CONTEXT,
+          traceState,
+        })
+      ).traceState,
+      traceState
+    );
+  });
 });
diff --git a/packages/opentelemetry-sdk-trace-base/test/common/sampler/TraceIdRatioBasedSampler.test.ts b/packages/opentelemetry-sdk-trace-base/test/common/sampler/TraceIdRatioBasedSampler.test.ts
@@ -18,11 +18,12 @@ import * as assert from 'assert';
 import * as api from '@opentelemetry/api';
 import { TraceIdRatioBasedSampler } from '../../../src/sampler/TraceIdRatioBasedSampler';
 
-const spanContext = (traceId = '1') =>
+const spanContext = (traceId = '1', traceState?: api.TraceState) =>
   api.trace.setSpanContext(api.ROOT_CONTEXT, {
     traceId,
     spanId: '1.1',
     traceFlags: api.TraceFlags.NONE,
+    traceState,
   });
 
 const traceId = (part: string) => ('0'.repeat(32) + part).slice(-32);
@@ -172,6 +173,16 @@ describe('TraceIdRatioBasedSampler', () => {
     );
   });
 
+  it('should forward the traceState', () => {
+    const sampler = new TraceIdRatioBasedSampler(1);
+    const traceState = api.createTraceState();
+    assert.strictEqual(
+      sampler.shouldSample(spanContext(traceId('1'), traceState), traceId('1'))
+        .traceState,
+      traceState
+    );
+  });
+
   it('should sample traces that a lower sampling ratio would sample', () => {
     const sampler10 = new TraceIdRatioBasedSampler(0.1);
     const sampler20 = new TraceIdRatioBasedSampler(0.2);