stack/apps/api/src/telemetry/span-context.service.spec.ts

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { SpanContextService } from "./span-context.service";
import { context, trace, type Span, type Context } from "@opentelemetry/api";
import { AsyncHooksContextManager } from "@opentelemetry/context-async-hooks";

describe("SpanContextService", () => {
  let service: SpanContextService;
  let contextManager: AsyncHooksContextManager;

  beforeEach(() => {
    // Set up context manager for proper context propagation in tests
    contextManager = new AsyncHooksContextManager();
    contextManager.enable();
    context.setGlobalContextManager(contextManager);
    service = new SpanContextService();
  });

  afterEach(() => {
    // Clean up context manager
    contextManager.disable();
    context.disable();
  });

  describe("getActiveSpan", () => {
    it("should return undefined when no span is active", () => {
      const activeSpan = service.getActiveSpan();
      expect(activeSpan).toBeUndefined();
    });

    it("should return the active span when one exists", () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("test-span");

      const result = context.with(trace.setSpan(context.active(), span), () => {
        return service.getActiveSpan();
      });

      expect(result).toBeDefined();
      expect(result).toBe(span);
      span.end();
    });
  });

  describe("getContext", () => {
    it("should return the current context", () => {
      const ctx = service.getContext();
      expect(ctx).toBeDefined();
    });

    it("should return the active context", () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("test-span");

      const result = context.with(trace.setSpan(context.active(), span), () => {
        return service.getContext();
      });

      expect(result).toBeDefined();
      span.end();
    });
  });

  describe("with", () => {
    it("should execute function within the provided context", () => {
      const customContext = context.active();
      let executedInContext = false;

      const result = service.with(customContext, () => {
        executedInContext = true;
        return "test-result";
      });

      expect(executedInContext).toBe(true);
      expect(result).toBe("test-result");
    });

    it("should propagate the context to the function", () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("test-span");
      const spanContext = trace.setSpan(context.active(), span);

      const result = service.with(spanContext, () => {
        const activeSpan = trace.getActiveSpan();
        return activeSpan;
      });

      expect(result).toBe(span);
      span.end();
    });

    it("should handle exceptions in the function", () => {
      const customContext = context.active();

      expect(() => {
        service.with(customContext, () => {
          throw new Error("Test error");
        });
      }).toThrow("Test error");
    });

    it("should return the function result", () => {
      const customContext = context.active();
      const result = service.with(customContext, () => {
        return { data: "test", count: 42 };
      });

      expect(result).toEqual({ data: "test", count: 42 });
    });
  });

  describe("withActiveSpan", () => {
    it("should set span as active for function execution", () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("test-span");

      const result = service.withActiveSpan(span, () => {
        const activeSpan = trace.getActiveSpan();
        return activeSpan;
      });

      expect(result).toBe(span);
      span.end();
    });

    it("should return the function result", () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("test-span");

      const result = service.withActiveSpan(span, () => {
        return "executed-with-span";
      });

      expect(result).toBe("executed-with-span");
      span.end();
    });

    it("should handle async operations", async () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("test-span");

      const result = await service.withActiveSpan(span, async () => {
        return new Promise((resolve) => {
          setTimeout(() => resolve("async-result"), 10);
        });
      });

      expect(result).toBe("async-result");
      span.end();
    });

    it("should propagate exceptions", () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("test-span");

      expect(() => {
        service.withActiveSpan(span, () => {
          throw new Error("Span execution error");
        });
      }).toThrow("Span execution error");

      span.end();
    });

    it("should nest spans correctly", () => {
      const tracer = trace.getTracer("test-tracer");
      const parentSpan = tracer.startSpan("parent-span");
      const childSpan = tracer.startSpan("child-span");

      const result = service.withActiveSpan(parentSpan, () => {
        return service.withActiveSpan(childSpan, () => {
          const activeSpan = trace.getActiveSpan();
          return activeSpan;
        });
      });

      expect(result).toBe(childSpan);
      childSpan.end();
      parentSpan.end();
    });
  });

  describe("integration scenarios", () => {
    it("should support complex span context propagation", () => {
      const tracer = trace.getTracer("test-tracer");
      const span1 = tracer.startSpan("span-1");
      const span2 = tracer.startSpan("span-2");

      // Execute with span1 active
      const ctx1 = trace.setSpan(context.active(), span1);
      const result1 = service.with(ctx1, () => {
        return service.getActiveSpan();
      });

      // Execute with span2 active
      const ctx2 = trace.setSpan(context.active(), span2);
      const result2 = service.with(ctx2, () => {
        return service.getActiveSpan();
      });

      expect(result1).toBe(span1);
      expect(result2).toBe(span2);

      span1.end();
      span2.end();
    });

    it("should handle context isolation", () => {
      const tracer = trace.getTracer("test-tracer");
      const span = tracer.startSpan("isolated-span");

      // Execute with span active
      service.withActiveSpan(span, () => {
        const activeInside = service.getActiveSpan();
        expect(activeInside).toBe(span);
      });

      // Outside the context, span should not be active
      const activeOutside = service.getActiveSpan();
      expect(activeOutside).not.toBe(span);

      span.end();
    });
  });
});