autotel/src/sampling.test.ts

Write Once, Observe Anywhere

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { TraceFlags } from '@opentelemetry/api';
import type { Link, SpanContext } from '@opentelemetry/api';
import {
  RandomSampler,
  AlwaysSampler,
  NeverSampler,
  AdaptiveSampler,
  UserIdSampler,
  CompositeSampler,
  FeatureFlagSampler,
  createLinkFromHeaders,
  extractLinksFromBatch,
  samplingPresets,
  resolveSamplingPreset,
  type SamplingContext,
} from './sampling';
import { type ILogger } from './logger';

describe('Sampling', () => {
  let mockLogger: ILogger;
  let context: SamplingContext;

  beforeEach(() => {
    mockLogger = {
      info: vi.fn(),
      warn: vi.fn(),
      error: vi.fn(),
      debug: vi.fn(),
    };

    context = {
      operationName: 'test.operation',
      args: [{ userId: '123', email: 'test@example.com' }],
    };
  });

  describe('RandomSampler', () => {
    it('should throw error for invalid sample rates', () => {
      expect(() => new RandomSampler(-0.1)).toThrow();
      expect(() => new RandomSampler(1.1)).toThrow();
    });

    it('should sample at 100%', () => {
      const sampler = new RandomSampler(1);
      const results = Array.from({ length: 100 }, () =>
        sampler.shouldSample(context),
      );

      expect(results.every((r) => r === true)).toBe(true);
    });

    it('should never sample at 0%', () => {
      const sampler = new RandomSampler(0);
      const results = Array.from({ length: 100 }, () =>
        sampler.shouldSample(context),
      );

      expect(results.every((r) => r === false)).toBe(true);
    });

    it('should sample approximately at the specified rate', () => {
      const sampler = new RandomSampler(0.5);
      const results = Array.from({ length: 1000 }, () =>
        sampler.shouldSample(context),
      );

      const sampleCount = results.filter(Boolean).length;
      // Allow 20% margin of error — random sampling is inherently noisy
      expect(sampleCount).toBeGreaterThan(400);
      expect(sampleCount).toBeLessThan(600);
    });
  });

  describe('AlwaysSampler', () => {
    it('should always sample', () => {
      const sampler = new AlwaysSampler();
      const results = Array.from({ length: 100 }, () =>
        sampler.shouldSample(context),
      );

      expect(results.every((r) => r === true)).toBe(true);
    });
  });

  describe('NeverSampler', () => {
    it('should never sample', () => {
      const sampler = new NeverSampler();
      const results = Array.from({ length: 100 }, () =>
        sampler.shouldSample(context),
      );

      expect(results.every((r) => r === false)).toBe(true);
    });
  });

  describe('AdaptiveSampler', () => {
    it('should throw error for invalid baseline sample rate', () => {
      expect(
        () =>
          new AdaptiveSampler({
            baselineSampleRate: -0.1,
          }),
      ).toThrow();
    });

    it('should indicate it needs tail sampling', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0.1,
      });

      expect(sampler.needsTailSampling()).toBe(true);
    });

    it('should always create spans (optimistic sampling)', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0, // Even with 0% baseline
        logger: mockLogger,
      });

      // FIX: shouldSample now always returns true for tail sampling
      const result = sampler.shouldSample(context);
      expect(result).toBe(true);
    });

    it('should keep error traces even when baseline would drop them', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0, // 0% baseline sampling
        alwaysSampleErrors: true,
        logger: mockLogger,
      });

      // FIX: Span is created (shouldSample returns true)
      const shouldSample = sampler.shouldSample(context);
      expect(shouldSample).toBe(true);

      // Tail sampling keeps error traces
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: false,
        duration: 100,
        error: new Error('Test error'),
      });

      expect(shouldKeep).toBe(true);
      // Pino-native: (extra, message)
      expect(mockLogger.debug).toHaveBeenCalledWith(
        expect.objectContaining({
          operation: 'test.operation',
          error: 'Test error',
        }),
        'Adaptive sampling: Keeping error trace',
      );
    });

    it('should keep slow traces even when baseline would drop them', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0, // 0% baseline sampling
        slowThresholdMs: 1000,
        alwaysSampleSlow: true,
        logger: mockLogger,
      });

      // FIX: Span is created (shouldSample returns true)
      const shouldSample = sampler.shouldSample(context);
      expect(shouldSample).toBe(true);

      // Tail sampling keeps slow traces
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 1500, // > 1000ms threshold
      });

      expect(shouldKeep).toBe(true);
      // Pino-native: (extra, message)
      expect(mockLogger.debug).toHaveBeenCalledWith(
        expect.objectContaining({
          operation: 'test.operation',
          duration: 1500,
        }),
        'Adaptive sampling: Keeping slow trace',
      );
    });

    it('should drop fast successful traces when baseline sampling says no', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0, // 0% baseline
        slowThresholdMs: 1000,
        logger: mockLogger,
      });

      // Span is created optimistically
      const shouldSample = sampler.shouldSample(context);
      expect(shouldSample).toBe(true);

      // Tail sampling drops fast/successful traces
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100, // < 1000ms threshold
      });

      expect(shouldKeep).toBe(false);
    });

    it('should keep fast successful traces when baseline sampling says yes', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 1, // 100% baseline
        slowThresholdMs: 1000,
        logger: mockLogger,
      });

      const shouldSample = sampler.shouldSample(context);
      expect(shouldSample).toBe(true);

      // Baseline sampled it, so keep it
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(true);
    });

    it('should respect alwaysSampleErrors flag', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        alwaysSampleErrors: false, // Don't force-sample errors
        logger: mockLogger,
      });

      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: false,
        duration: 100,
        error: new Error('Test error'),
      });

      // With alwaysSampleErrors=false and baseline=0, errors are dropped
      expect(shouldKeep).toBe(false);
    });

    it('should respect alwaysSampleSlow flag', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        slowThresholdMs: 1000,
        alwaysSampleSlow: false, // Don't force-sample slow requests
        logger: mockLogger,
      });

      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 1500,
      });

      // With alwaysSampleSlow=false and baseline=0, slow requests are dropped
      expect(shouldKeep).toBe(false);
    });
  });

  describe('UserIdSampler', () => {
    const extractUserId = (args: unknown[]) => {
      const firstArg = args[0] as { userId?: string };
      return firstArg?.userId;
    };

    it('should always sample specific users', () => {
      const sampler = new UserIdSampler({
        baselineSampleRate: 0,
        alwaysSampleUsers: ['vip_123'],
        extractUserId,
        logger: mockLogger,
      });

      const vipContext: SamplingContext = {
        operationName: 'test.operation',
        args: [{ userId: 'vip_123' }],
      };

      expect(sampler.shouldSample(vipContext)).toBe(true);
      // Pino-native: (extra, message)
      expect(mockLogger.debug).toHaveBeenCalledWith(
        {
          operation: 'test.operation',
          userId: 'vip_123',
        },
        'Sampling user request',
      );
    });

    it('should use consistent per-user sampling', () => {
      const sampler = new UserIdSampler({
        baselineSampleRate: 0.5,
        extractUserId,
        logger: mockLogger,
      });

      const user123Context: SamplingContext = {
        operationName: 'test.operation',
        args: [{ userId: 'user_123' }],
      };

      // Same user should always get same result
      const result1 = sampler.shouldSample(user123Context);
      const result2 = sampler.shouldSample(user123Context);
      const result3 = sampler.shouldSample(user123Context);

      expect(result1).toBe(result2);
      expect(result2).toBe(result3);
    });

    it('should add and remove users from always-sample list', () => {
      const sampler = new UserIdSampler({
        baselineSampleRate: 0,
        extractUserId,
      });

      sampler.addAlwaysSampleUsers('user_1', 'user_2');

      const user1Context: SamplingContext = {
        operationName: 'test.operation',
        args: [{ userId: 'user_1' }],
      };

      expect(sampler.shouldSample(user1Context)).toBe(true);

      sampler.removeAlwaysSampleUsers('user_1');
      expect(sampler.shouldSample(user1Context)).toBe(false);
    });

    it('should fallback to random sampling when no user ID', () => {
      const sampler = new UserIdSampler({
        baselineSampleRate: 1,
        extractUserId: (args) => (args[0] as { userId?: string })?.userId,
      });

      const noUserContext: SamplingContext = {
        operationName: 'test.operation',
        args: [{}],
      };

      expect(sampler.shouldSample(noUserContext)).toBe(true);
    });
  });

  describe('CompositeSampler', () => {
    it('should throw error with no child samplers', () => {
      expect(() => new CompositeSampler([])).toThrow();
    });

    it('should sample if any child sampler returns true', () => {
      const sampler = new CompositeSampler([
        new NeverSampler(),
        new AlwaysSampler(),
        new NeverSampler(),
      ]);

      expect(sampler.shouldSample(context)).toBe(true);
    });

    it('should not sample if all child samplers return false', () => {
      const sampler = new CompositeSampler([
        new NeverSampler(),
        new NeverSampler(),
      ]);

      expect(sampler.shouldSample(context)).toBe(false);
    });
  });

  describe('FeatureFlagSampler', () => {
    const extractFlags = (
      args: unknown[],
      metadata?: Record<string, unknown>,
    ) => {
      const firstArg = args[0] as { flags?: string[] };
      return firstArg?.flags || (metadata?.featureFlags as string[]);
    };

    it('should always sample requests with monitored flags', () => {
      const sampler = new FeatureFlagSampler({
        baselineSampleRate: 0,
        alwaysSampleFlags: ['new_checkout', 'experimental_ui'],
        extractFlags,
        logger: mockLogger,
      });

      const flagContext: SamplingContext = {
        operationName: 'test.operation',
        args: [{ flags: ['new_checkout'] }],
      };

      expect(sampler.shouldSample(flagContext)).toBe(true);
      // Pino-native: (extra, message)
      expect(mockLogger.debug).toHaveBeenCalledWith(
        {
          operation: 'test.operation',
          flags: ['new_checkout'],
        },
        'Sampling feature flag request',
      );
    });

    it('should use baseline sampling for non-monitored flags', () => {
      const sampler = new FeatureFlagSampler({
        baselineSampleRate: 0,
        alwaysSampleFlags: ['monitored_flag'],
        extractFlags,
      });

      const flagContext: SamplingContext = {
        operationName: 'test.operation',
        args: [{ flags: ['other_flag'] }],
      };

      expect(sampler.shouldSample(flagContext)).toBe(false);
    });

    it('should add and remove flags', () => {
      const sampler = new FeatureFlagSampler({
        baselineSampleRate: 0,
        extractFlags,
      });

      sampler.addAlwaysSampleFlags('flag_1', 'flag_2');

      const flag1Context: SamplingContext = {
        operationName: 'test.operation',
        args: [{ flags: ['flag_1'] }],
      };

      expect(sampler.shouldSample(flag1Context)).toBe(true);

      sampler.removeAlwaysSampleFlags('flag_1');
      expect(sampler.shouldSample(flag1Context)).toBe(false);
    });
  });

  describe('Real-world QA in Production scenarios', () => {
    it('should always capture failed email deliveries from article example', () => {
      // From article: "We set up another alert that let us know if our
      // email-sending microservice was unable to process a request"
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0.1, // 10% baseline
        alwaysSampleErrors: true, // Always capture failures
      });

      const emailContext: SamplingContext = {
        operationName: 'email.send',
        args: [{ to: 'school@example.com' }],
      };

      sampler.shouldSample(emailContext);

      // Email fails due to invalid address
      const shouldKeep = sampler.shouldKeepTrace(emailContext, {
        success: false,
        duration: 100,
        error: new Error('Invalid email address'),
      });

      expect(shouldKeep).toBe(true);
    });

    it('should always trace slow job applications from article example', () => {
      // From article: Monitor if teachers are able to submit applications
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0.05, // 5% baseline
        slowThresholdMs: 2000, // Slow if > 2s
        alwaysSampleSlow: true,
      });

      const applicationContext: SamplingContext = {
        operationName: 'application.submit',
        args: [{ jobId: '123', teacherId: '456' }],
      };

      sampler.shouldSample(applicationContext);

      // Application takes too long
      const shouldKeep = sampler.shouldKeepTrace(applicationContext, {
        success: true,
        duration: 3000, // > 2s threshold
      });

      expect(shouldKeep).toBe(true);
    });

    it('should always trace VIP users', () => {
      const extractUserId = (args: unknown[]) => {
        const firstArg = args[0] as { userId?: string };
        return firstArg?.userId;
      };

      const sampler = new UserIdSampler({
        baselineSampleRate: 0.01, // 1% of normal users
        alwaysSampleUsers: ['vip_school_123'], // Always trace VIP schools
        extractUserId,
      });

      const vipContext: SamplingContext = {
        operationName: 'application.receive',
        args: [{ userId: 'vip_school_123' }],
      };

      expect(sampler.shouldSample(vipContext)).toBe(true);
    });

    it('should always trace A/B test variants for correlation', () => {
      const extractFlags = (args: unknown[]) => {
        const firstArg = args[0] as { experimentFlags?: string[] };
        return firstArg?.experimentFlags;
      };

      const sampler = new FeatureFlagSampler({
        baselineSampleRate: 0.05,
        alwaysSampleFlags: ['new_application_form'], // Always trace experiment
        extractFlags,
      });

      const experimentContext: SamplingContext = {
        operationName: 'application.submit',
        args: [{ experimentFlags: ['new_application_form'] }],
      };

      expect(sampler.shouldSample(experimentContext)).toBe(true);
    });
  });

  describe('AdaptiveSampler - Links-based sampling', () => {
    let mockLogger: ILogger;

    beforeEach(() => {
      mockLogger = {
        info: vi.fn(),
        warn: vi.fn(),
        error: vi.fn(),
        debug: vi.fn(),
      };
    });

    // Helper to create a SpanContext
    const createSpanContext = (sampled: boolean): SpanContext => ({
      traceId: '0af7651916cd43dd8448eb211c80319c',
      spanId: 'b7ad6b7169203331',
      traceFlags: sampled ? TraceFlags.SAMPLED : TraceFlags.NONE,
      isRemote: true,
    });

    // Helper to create a Link
    const createLink = (
      sampled: boolean,
      attributes?: Record<string, unknown>,
    ): Link => ({
      context: createSpanContext(sampled),
      attributes: attributes ?? {},
    });

    it('should have linksBased disabled by default', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0.1,
      });

      const context: SamplingContext = {
        operationName: 'test.operation',
        args: [],
        links: [createLink(true)], // Sampled link
      };

      sampler.shouldSample(context);

      // With linksBased=false (default), sampled links don't affect decision
      // With baselineSampleRate=0.1 and random chance, we can't deterministically test
      // So we test with 0% baseline - the link should NOT cause it to be kept
      const samplerStrict = new AdaptiveSampler({
        baselineSampleRate: 0,
        linksBased: false, // Explicitly disabled
      });

      samplerStrict.shouldSample(context);
      const shouldKeep = samplerStrict.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(false);
    });

    it('should throw error for invalid linksRate', () => {
      expect(
        () =>
          new AdaptiveSampler({
            linksRate: -0.1,
          }),
      ).toThrow('Links rate must be between 0 and 1');

      expect(
        () =>
          new AdaptiveSampler({
            linksRate: 1.5,
          }),
      ).toThrow('Links rate must be between 0 and 1');
    });

    it('should keep traces with sampled links when linksBased=true', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0, // 0% baseline
        linksBased: true,
        linksRate: 1, // 100% of linked spans kept
        logger: mockLogger,
      });

      const context: SamplingContext = {
        operationName: 'consumer.process',
        args: [],
        links: [createLink(true)], // Sampled link
      };

      sampler.shouldSample(context);
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(true);
      // Pino-native: (extra, message)
      expect(mockLogger.debug).toHaveBeenCalledWith(
        expect.objectContaining({
          operation: 'consumer.process',
          linkCount: 1,
        }),
        'Adaptive sampling: Keeping trace due to sampled link',
      );
    });

    it('should drop traces with unsampled links when linksBased=true', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        linksBased: true,
        linksRate: 1,
      });

      const context: SamplingContext = {
        operationName: 'consumer.process',
        args: [],
        links: [createLink(false)], // NOT sampled
      };

      sampler.shouldSample(context);
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(false);
    });

    it('should keep trace if ANY link is sampled (fan-in)', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        linksBased: true,
        linksRate: 1,
      });

      const context: SamplingContext = {
        operationName: 'batch.process',
        args: [],
        links: [
          createLink(false), // Not sampled
          createLink(true), // Sampled
          createLink(false), // Not sampled
        ],
      };

      sampler.shouldSample(context);
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(true);
    });

    it('should respect linksRate for probabilistic link sampling', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        linksBased: true,
        linksRate: 0, // 0% - never keep linked spans
      });

      const context: SamplingContext = {
        operationName: 'consumer.process',
        args: [],
        links: [createLink(true)],
      };

      sampler.shouldSample(context);
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(false);
    });

    it('should prioritize errors over links-based sampling', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        linksBased: true,
        linksRate: 0, // Would drop linked spans
        alwaysSampleErrors: true,
      });

      const context: SamplingContext = {
        operationName: 'consumer.process',
        args: [],
        links: [createLink(false)], // No sampled links
      };

      sampler.shouldSample(context);
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: false, // Error!
        duration: 100,
        error: new Error('Processing failed'),
      });

      expect(shouldKeep).toBe(true); // Errors always kept
    });

    it('should handle empty links array', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        linksBased: true,
        linksRate: 1,
      });

      const context: SamplingContext = {
        operationName: 'test.operation',
        args: [],
        links: [],
      };

      sampler.shouldSample(context);
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(false); // No links, baseline=0
    });

    it('should handle undefined links', () => {
      const sampler = new AdaptiveSampler({
        baselineSampleRate: 0,
        linksBased: true,
        linksRate: 1,
      });

      const context: SamplingContext = {
        operationName: 'test.operation',
        args: [],
        // links is undefined
      };

      sampler.shouldSample(context);
      const shouldKeep = sampler.shouldKeepTrace(context, {
        success: true,
        duration: 100,
      });

      expect(shouldKeep).toBe(false);
    });

    describe('hasSampledLink helper', () => {
      it('should return false for empty links', () => {
        const sampler = new AdaptiveSampler({ linksBased: true });
        expect(sampler.hasSampledLink([])).toBe(false);
      });

      it('should return true if any link is sampled', () => {
        const sampler = new AdaptiveSampler({ linksBased: true });
        const links = [createLink(false), createLink(true)];
        expect(sampler.hasSampledLink(links)).toBe(true);
      });

      it('should return false if no links are sampled', () => {
        const sampler = new AdaptiveSampler({ linksBased: true });
        const links = [createLink(false), createLink(false)];
        expect(sampler.hasSampledLink(links)).toBe(false);
      });
    });
  });

  describe('Link Helper Functions', () => {
    describe('createLinkFromHeaders', () => {
      it('should create link from valid W3C traceparent header', () => {
        const headers = {
          traceparent:
            '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
        };

        const link = createLinkFromHeaders(headers);

        expect(link).not.toBeNull();
        expect(link?.context.traceId).toBe('0af7651916cd43dd8448eb211c80319c');
        expect(link?.context.spanId).toBe('b7ad6b7169203331');
        expect(link?.context.traceFlags).toBe(TraceFlags.SAMPLED);
      });

      it('should create link from unsampled traceparent', () => {
        const headers = {
          traceparent:
            '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00',
        };

        const link = createLinkFromHeaders(headers);

        expect(link).not.toBeNull();
        expect(link?.context.traceFlags).toBe(TraceFlags.NONE);
      });

      it('should include custom attributes in link', () => {
        const headers = {
          traceparent:
            '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
        };
        const attributes = {
          'messaging.system': 'kafka',
          'custom.attr': 'value',
        };

        const link = createLinkFromHeaders(headers, attributes);

        expect(link?.attributes).toEqual(attributes);
      });

      it('should return null for invalid/missing traceparent', () => {
        expect(createLinkFromHeaders({})).toBeNull();
        expect(createLinkFromHeaders({ traceparent: 'invalid' })).toBeNull();
        expect(createLinkFromHeaders({ traceparent: '' })).toBeNull();
      });

      it('should return null for all-zero trace IDs', () => {
        // All zeros is an invalid trace context
        const headers = {
          traceparent:
            '00-00000000000000000000000000000000-0000000000000000-01',
        };

        const link = createLinkFromHeaders(headers);
        expect(link).toBeNull();
      });
    });

    describe('extractLinksFromBatch', () => {
      it('should extract links from batch of messages', () => {
        const messages = [
          {
            body: 'message1',
            headers: {
              traceparent:
                '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
            },
          },
          {
            body: 'message2',
            headers: {
              traceparent:
                '00-1af7651916cd43dd8448eb211c80319c-c7ad6b7169203331-01',
            },
          },
        ];

        const links = extractLinksFromBatch(messages);

        expect(links).toHaveLength(2);
        expect(links[0].context.traceId).toBe(
          '0af7651916cd43dd8448eb211c80319c',
        );
        expect(links[1].context.traceId).toBe(
          '1af7651916cd43dd8448eb211c80319c',
        );
      });

      it('should add message index as link attribute', () => {
        const messages = [
          {
            body: 'message1',
            headers: {
              traceparent:
                '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
            },
          },
          {
            body: 'message2',
            headers: {
              traceparent:
                '00-1af7651916cd43dd8448eb211c80319c-c7ad6b7169203331-01',
            },
          },
        ];

        const links = extractLinksFromBatch(messages);

        expect(links[0].attributes?.['messaging.batch.message_index']).toBe(0);
        expect(links[1].attributes?.['messaging.batch.message_index']).toBe(1);
      });

      it('should use custom headers key', () => {
        const messages = [
          {
            body: 'message1',
            metadata: {
              traceparent:
                '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
            },
          },
        ];

        const links = extractLinksFromBatch(messages, 'metadata');

        expect(links).toHaveLength(1);
      });

      it('should skip messages without headers', () => {
        const messages = [
          { body: 'message1' }, // No headers
          {
            body: 'message2',
            headers: {
              traceparent:
                '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
            },
          },
          { body: 'message3', headers: null }, // Null headers
        ];

        const links = extractLinksFromBatch(messages);

        expect(links).toHaveLength(1);
      });

      it('should skip messages with invalid trace context', () => {
        const messages = [
          {
            body: 'message1',
            headers: { traceparent: 'invalid-traceparent' },
          },
          {
            body: 'message2',
            headers: {
              traceparent:
                '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01',
            },
          },
        ];

        const links = extractLinksFromBatch(messages);

        expect(links).toHaveLength(1);
        expect(links[0].context.traceId).toBe(
          '0af7651916cd43dd8448eb211c80319c',
        );
      });

      it('should return empty array for empty batch', () => {
        const links = extractLinksFromBatch([]);
        expect(links).toEqual([]);
      });
    });
  });

  describe('samplingPresets', () => {
    it('development() returns AlwaysSampler', () => {
      const sampler = samplingPresets.development();
      expect(sampler).toBeInstanceOf(AlwaysSampler);
      expect(sampler.shouldSample(context)).toBe(true);
    });

    it('errorsOnly() returns AdaptiveSampler that drops healthy baseline', () => {
      const sampler = samplingPresets.errorsOnly();
      expect(sampler).toBeInstanceOf(AdaptiveSampler);
      sampler.shouldSample(context); // prime the WeakMap baseline decision
      // Baseline is 0, so shouldKeepTrace for successful fast requests = false
      expect(
        sampler.shouldKeepTrace!(context, { success: true, duration: 50 }),
      ).toBe(false);
    });

    it('errorsOnly() keeps errors', () => {
      const sampler = samplingPresets.errorsOnly();
      sampler.shouldSample(context); // prime the baseline decision
      expect(
        sampler.shouldKeepTrace!(context, {
          success: false,
          duration: 50,
          error: new Error('fail'),
        }),
      ).toBe(true);
    });

    it('production() returns AdaptiveSampler with 10% baseline', () => {
      const sampler = samplingPresets.production();
      expect(sampler).toBeInstanceOf(AdaptiveSampler);
      expect(sampler.needsTailSampling!()).toBe(true);
    });

    it('production() keeps errors', () => {
      const sampler = samplingPresets.production();
      sampler.shouldSample(context);
      expect(
        sampler.shouldKeepTrace!(context, {
          success: false,
          duration: 50,
          error: new Error('fail'),
        }),
      ).toBe(true);
    });

    it('production() accepts overrides', () => {
      const sampler = samplingPresets.production({ baselineSampleRate: 1.0 });
      sampler.shouldSample(context);
      // With 100% baseline, all healthy traffic is kept
      expect(
        sampler.shouldKeepTrace!(context, { success: true, duration: 50 }),
      ).toBe(true);
    });

    it('off() returns NeverSampler', () => {
      const sampler = samplingPresets.off();
      expect(sampler).toBeInstanceOf(NeverSampler);
      expect(sampler.shouldSample(context)).toBe(false);
    });
  });

  describe('resolveSamplingPreset', () => {
    it('resolves development', () => {
      const sampler = resolveSamplingPreset('development');
      expect(sampler).toBeInstanceOf(AlwaysSampler);
    });

    it('resolves errors-only', () => {
      const sampler = resolveSamplingPreset('errors-only');
      expect(sampler).toBeInstanceOf(AdaptiveSampler);
    });

    it('resolves production', () => {
      const sampler = resolveSamplingPreset('production');
      expect(sampler).toBeInstanceOf(AdaptiveSampler);
    });

    it('resolves off', () => {
      const sampler = resolveSamplingPreset('off');
      expect(sampler).toBeInstanceOf(NeverSampler);
    });

    it('throws on invalid preset with helpful message', () => {
      expect(() => resolveSamplingPreset('banana' as any)).toThrow(
        /Unknown sampling preset: "banana".*Valid presets: development, errors-only, production, off/,
      );
    });
  });
});