stack/apps/api/src/mosaic-telemetry/mosaic-telemetry.service.spec.ts

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { ConfigService } from "@nestjs/config";
import { MOSAIC_TELEMETRY_ENV } from "./mosaic-telemetry.config";
import type {
  TaskCompletionEvent,
  PredictionQuery,
  PredictionResponse,
} from "@mosaicstack/telemetry-client";
import { TaskType, Complexity, Provider, Outcome } from "@mosaicstack/telemetry-client";

// Track mock instances created during tests
const mockStartFn = vi.fn();
const mockStopFn = vi.fn().mockResolvedValue(undefined);
const mockTrackFn = vi.fn();
const mockGetPredictionFn = vi.fn().mockReturnValue(null);
const mockRefreshPredictionsFn = vi.fn().mockResolvedValue(undefined);
const mockBuildFn = vi.fn().mockReturnValue({ event_id: "test-event-id" });

vi.mock("@mosaicstack/telemetry-client", async (importOriginal) => {
  const actual = await importOriginal<typeof import("@mosaicstack/telemetry-client")>();

  class MockTelemetryClient {
    private _isRunning = false;

    constructor(_config: unknown) {
      // no-op
    }

    get eventBuilder() {
      return { build: mockBuildFn };
    }

    start(): void {
      this._isRunning = true;
      mockStartFn();
    }

    async stop(): Promise<void> {
      this._isRunning = false;
      await mockStopFn();
    }

    track(event: unknown): void {
      mockTrackFn(event);
    }

    getPrediction(query: unknown): unknown {
      return mockGetPredictionFn(query);
    }

    async refreshPredictions(queries: unknown): Promise<void> {
      await mockRefreshPredictionsFn(queries);
    }

    get queueSize(): number {
      return 0;
    }

    get isRunning(): boolean {
      return this._isRunning;
    }
  }

  return {
    ...actual,
    TelemetryClient: MockTelemetryClient,
  };
});

// Lazy-import the service after the mock is in place
const { MosaicTelemetryService } = await import("./mosaic-telemetry.service");

/**
 * Create a ConfigService mock that returns environment values from the provided map.
 */
function createConfigService(
  envMap: Record<string, string | undefined> = {},
): ConfigService {
  const configService = {
    get: vi.fn((key: string, defaultValue?: string): string => {
      const value = envMap[key];
      if (value !== undefined) {
        return value;
      }
      return defaultValue ?? "";
    }),
  } as unknown as ConfigService;
  return configService;
}

/**
 * Default env config for an enabled telemetry service.
 */
const ENABLED_CONFIG: Record<string, string> = {
  [MOSAIC_TELEMETRY_ENV.ENABLED]: "true",
  [MOSAIC_TELEMETRY_ENV.SERVER_URL]: "https://tel.test.local",
  [MOSAIC_TELEMETRY_ENV.API_KEY]: "a".repeat(64),
  [MOSAIC_TELEMETRY_ENV.INSTANCE_ID]: "550e8400-e29b-41d4-a716-446655440000",
  [MOSAIC_TELEMETRY_ENV.DRY_RUN]: "false",
};

/**
 * Create a minimal TaskCompletionEvent for testing.
 */
function createTestEvent(): TaskCompletionEvent {
  return {
    schema_version: "1.0.0",
    event_id: "test-event-123",
    timestamp: new Date().toISOString(),
    instance_id: "550e8400-e29b-41d4-a716-446655440000",
    task_duration_ms: 5000,
    task_type: TaskType.FEATURE,
    complexity: Complexity.MEDIUM,
    harness: "claude-code" as TaskCompletionEvent["harness"],
    model: "claude-sonnet-4-20250514",
    provider: Provider.ANTHROPIC,
    estimated_input_tokens: 1000,
    estimated_output_tokens: 500,
    actual_input_tokens: 1100,
    actual_output_tokens: 450,
    estimated_cost_usd_micros: 5000,
    actual_cost_usd_micros: 4800,
    quality_gate_passed: true,
    quality_gates_run: [],
    quality_gates_failed: [],
    context_compactions: 0,
    context_rotations: 0,
    context_utilization_final: 0.45,
    outcome: Outcome.SUCCESS,
    retry_count: 0,
  };
}

describe("MosaicTelemetryService", () => {
  let service: InstanceType<typeof MosaicTelemetryService>;

  afterEach(async () => {
    if (service) {
      await service.onModuleDestroy();
    }
    vi.clearAllMocks();
  });

  describe("onModuleInit", () => {
    it("should initialize the client when enabled with valid config", () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);

      service.onModuleInit();

      expect(mockStartFn).toHaveBeenCalledOnce();
      expect(service.isEnabled).toBe(true);
    });

    it("should not initialize client when disabled", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);

      service.onModuleInit();

      expect(mockStartFn).not.toHaveBeenCalled();
      expect(service.isEnabled).toBe(false);
    });

    it("should disable when server URL is missing", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.SERVER_URL]: "",
      });
      service = new MosaicTelemetryService(configService);

      service.onModuleInit();

      expect(service.isEnabled).toBe(false);
    });

    it("should disable when API key is missing", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.API_KEY]: "",
      });
      service = new MosaicTelemetryService(configService);

      service.onModuleInit();

      expect(service.isEnabled).toBe(false);
    });

    it("should disable when instance ID is missing", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.INSTANCE_ID]: "",
      });
      service = new MosaicTelemetryService(configService);

      service.onModuleInit();

      expect(service.isEnabled).toBe(false);
    });

    it("should log dry-run mode when configured", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.DRY_RUN]: "true",
      });
      service = new MosaicTelemetryService(configService);

      service.onModuleInit();

      expect(mockStartFn).toHaveBeenCalledOnce();
    });
  });

  describe("onModuleDestroy", () => {
    it("should stop the client on shutdown", async () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      await service.onModuleDestroy();

      expect(mockStopFn).toHaveBeenCalledOnce();
    });

    it("should not throw when client is not initialized (disabled)", async () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      await expect(service.onModuleDestroy()).resolves.not.toThrow();
    });

    it("should not throw when called multiple times", async () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      await service.onModuleDestroy();
      await expect(service.onModuleDestroy()).resolves.not.toThrow();
    });
  });

  describe("trackTaskCompletion", () => {
    it("should queue event via client.track() when enabled", () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      const event = createTestEvent();
      service.trackTaskCompletion(event);

      expect(mockTrackFn).toHaveBeenCalledWith(event);
    });

    it("should be a no-op when disabled", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      const event = createTestEvent();
      service.trackTaskCompletion(event);

      expect(mockTrackFn).not.toHaveBeenCalled();
    });
  });

  describe("getPrediction", () => {
    const testQuery: PredictionQuery = {
      task_type: TaskType.FEATURE,
      model: "claude-sonnet-4-20250514",
      provider: Provider.ANTHROPIC,
      complexity: Complexity.MEDIUM,
    };

    it("should return cached prediction when available", () => {
      const mockPrediction: PredictionResponse = {
        prediction: {
          input_tokens: { p10: 100, p25: 200, median: 300, p75: 400, p90: 500 },
          output_tokens: { p10: 50, p25: 100, median: 150, p75: 200, p90: 250 },
          cost_usd_micros: { median: 5000 },
          duration_ms: { median: 10000 },
          correction_factors: { input: 1.0, output: 1.0 },
          quality: { gate_pass_rate: 0.95, success_rate: 0.90 },
        },
        metadata: {
          sample_size: 100,
          fallback_level: 0,
          confidence: "high",
          last_updated: new Date().toISOString(),
          cache_hit: true,
        },
      };

      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      mockGetPredictionFn.mockReturnValueOnce(mockPrediction);

      const result = service.getPrediction(testQuery);

      expect(result).toEqual(mockPrediction);
      expect(mockGetPredictionFn).toHaveBeenCalledWith(testQuery);
    });

    it("should return null when disabled", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      const result = service.getPrediction(testQuery);

      expect(result).toBeNull();
    });

    it("should return null when no cached prediction exists", () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      mockGetPredictionFn.mockReturnValueOnce(null);

      const result = service.getPrediction(testQuery);

      expect(result).toBeNull();
    });
  });

  describe("refreshPredictions", () => {
    const testQueries: PredictionQuery[] = [
      {
        task_type: TaskType.FEATURE,
        model: "claude-sonnet-4-20250514",
        provider: Provider.ANTHROPIC,
        complexity: Complexity.MEDIUM,
      },
    ];

    it("should call client.refreshPredictions when enabled", async () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      await service.refreshPredictions(testQueries);

      expect(mockRefreshPredictionsFn).toHaveBeenCalledWith(testQueries);
    });

    it("should be a no-op when disabled", async () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      await service.refreshPredictions(testQueries);

      expect(mockRefreshPredictionsFn).not.toHaveBeenCalled();
    });
  });

  describe("eventBuilder", () => {
    it("should return EventBuilder when enabled", () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      const builder = service.eventBuilder;

      expect(builder).toBeDefined();
      expect(builder).not.toBeNull();
      expect(typeof builder?.build).toBe("function");
    });

    it("should return null when disabled", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      const builder = service.eventBuilder;

      expect(builder).toBeNull();
    });
  });

  describe("isEnabled", () => {
    it("should return true when client is running", () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      expect(service.isEnabled).toBe(true);
    });

    it("should return false when disabled", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      expect(service.isEnabled).toBe(false);
    });
  });

  describe("queueSize", () => {
    it("should return 0 when disabled", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      expect(service.queueSize).toBe(0);
    });

    it("should delegate to client.queueSize when enabled", () => {
      const configService = createConfigService(ENABLED_CONFIG);
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      expect(service.queueSize).toBe(0);
    });
  });

  describe("disabled mode (comprehensive)", () => {
    beforeEach(() => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.ENABLED]: "false",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();
    });

    it("should not make any HTTP calls when disabled", () => {
      const event = createTestEvent();
      service.trackTaskCompletion(event);

      expect(mockTrackFn).not.toHaveBeenCalled();
      expect(mockStartFn).not.toHaveBeenCalled();
    });

    it("should safely handle all method calls when disabled", async () => {
      expect(() => service.trackTaskCompletion(createTestEvent())).not.toThrow();
      expect(
        service.getPrediction({
          task_type: TaskType.FEATURE,
          model: "test",
          provider: Provider.ANTHROPIC,
          complexity: Complexity.LOW,
        }),
      ).toBeNull();
      await expect(service.refreshPredictions([])).resolves.not.toThrow();
      expect(service.eventBuilder).toBeNull();
      expect(service.isEnabled).toBe(false);
      expect(service.queueSize).toBe(0);
    });
  });

  describe("dry-run mode", () => {
    it("should create client in dry-run mode", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.DRY_RUN]: "true",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      expect(mockStartFn).toHaveBeenCalledOnce();
      expect(service.isEnabled).toBe(true);
    });

    it("should accept events in dry-run mode", () => {
      const configService = createConfigService({
        ...ENABLED_CONFIG,
        [MOSAIC_TELEMETRY_ENV.DRY_RUN]: "true",
      });
      service = new MosaicTelemetryService(configService);
      service.onModuleInit();

      const event = createTestEvent();
      service.trackTaskCompletion(event);

      expect(mockTrackFn).toHaveBeenCalledWith(event);
    });
  });
});