stack/apps/api/src/llm/llm-manager.service.spec.ts

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { LlmManagerService } from "./llm-manager.service";
import { PrismaService } from "../prisma/prisma.service";
import type { LlmProviderInstance } from "@prisma/client";
import type {
  LlmProviderInterface,
  LlmProviderHealthStatus,
} from "./providers/llm-provider.interface";

// Mock the OllamaProvider module before importing LlmManagerService
vi.mock("./providers/ollama.provider", () => {
  class MockOllamaProvider {
    readonly name = "Ollama";
    readonly type = "ollama" as const;
    private config: { endpoint: string; timeout?: number };

    constructor(config: { endpoint: string; timeout?: number }) {
      this.config = config;
    }

    async initialize(): Promise<void> {
      // No-op for testing
    }

    async checkHealth() {
      return {
        healthy: true,
        provider: "ollama",
        endpoint: this.config.endpoint,
      };
    }

    async listModels(): Promise<string[]> {
      return ["model1", "model2"];
    }

    async chat() {
      return {
        model: "test",
        message: { role: "assistant", content: "Mock response" },
        done: true,
      };
    }

    async *chatStream() {
      yield {
        model: "test",
        message: { role: "assistant", content: "Mock stream" },
        done: true,
      };
    }

    async embed() {
      return {
        model: "test",
        embeddings: [[0.1, 0.2, 0.3]],
      };
    }

    getConfig() {
      return { ...this.config };
    }
  }

  return {
    OllamaProvider: MockOllamaProvider,
  };
});

/**
 * Mock provider for testing purposes
 */
class MockProvider implements LlmProviderInterface {
  readonly name = "MockProvider";
  readonly type = "ollama" as const;

  constructor(private config: { endpoint: string }) {}

  async initialize(): Promise<void> {
    // No-op for testing
  }

  async checkHealth(): Promise<LlmProviderHealthStatus> {
    return {
      healthy: true,
      provider: "ollama",
      endpoint: this.config.endpoint,
    };
  }

  async listModels(): Promise<string[]> {
    return ["model1", "model2"];
  }

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  async chat(request: any): Promise<any> {
    return {
      model: request.model,
      message: { role: "assistant", content: "Mock response" },
      done: true,
    };
  }

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  async *chatStream(request: any): AsyncGenerator<any> {
    yield {
      model: request.model,
      message: { role: "assistant", content: "Mock stream" },
      done: true,
    };
  }

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  async embed(request: any): Promise<any> {
    return {
      model: request.model,
      embeddings: [[0.1, 0.2, 0.3]],
    };
  }

  getConfig(): { endpoint: string } {
    return { ...this.config };
  }
}

/**
 * Unhealthy mock provider for testing error scenarios
 */
class UnhealthyMockProvider extends MockProvider {
  async checkHealth(): Promise<LlmProviderHealthStatus> {
    return {
      healthy: false,
      provider: "ollama",
      endpoint: this.getConfig().endpoint,
      error: "Connection failed",
    };
  }
}

describe("LlmManagerService", () => {
  let service: LlmManagerService;
  let prisma: PrismaService;

  const mockProviderInstance: LlmProviderInstance = {
    id: "550e8400-e29b-41d4-a716-446655440000",
    providerType: "ollama",
    displayName: "Test Ollama",
    userId: null,
    config: { endpoint: "http://localhost:11434", timeout: 30000 },
    isDefault: true,
    isEnabled: true,
    createdAt: new Date(),
    updatedAt: new Date(),
  };

  const mockUserProviderInstance: LlmProviderInstance = {
    id: "550e8400-e29b-41d4-a716-446655440001",
    providerType: "ollama",
    displayName: "User Ollama",
    userId: "user-123",
    config: { endpoint: "http://user-ollama:11434", timeout: 30000 },
    isDefault: false,
    isEnabled: true,
    createdAt: new Date(),
    updatedAt: new Date(),
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        LlmManagerService,
        {
          provide: PrismaService,
          useValue: {
            llmProviderInstance: {
              findMany: vi.fn(),
              findUnique: vi.fn(),
              findFirst: vi.fn(),
            },
          },
        },
      ],
    }).compile();

    service = module.get<LlmManagerService>(LlmManagerService);
    prisma = module.get<PrismaService>(PrismaService);
  });

  afterEach(() => {
    vi.clearAllMocks();
  });

  describe("initialization", () => {
    it("should be defined", () => {
      expect(service).toBeDefined();
    });

    it("should load enabled providers from database on module init", async () => {
      const findManySpy = vi
        .spyOn(prisma.llmProviderInstance, "findMany")
        .mockResolvedValue([mockProviderInstance]);

      await service.onModuleInit();

      expect(findManySpy).toHaveBeenCalledWith({
        where: { isEnabled: true },
      });
    });

    it("should handle empty database gracefully", async () => {
      vi.spyOn(prisma.llmProviderInstance, "findMany").mockResolvedValue([]);

      await service.onModuleInit();

      const providers = await service.getAllProviders();
      expect(providers).toHaveLength(0);
    });

    it("should skip disabled providers during initialization", async () => {
      // Mock the database to return only enabled providers (empty array since all are disabled)
      vi.spyOn(prisma.llmProviderInstance, "findMany").mockResolvedValue([]);

      await service.onModuleInit();

      const providers = await service.getAllProviders();
      expect(providers).toHaveLength(0);
    });

    it("should initialize all loaded providers", async () => {
      vi.spyOn(prisma.llmProviderInstance, "findMany").mockResolvedValue([mockProviderInstance]);

      await service.onModuleInit();

      const provider = await service.getProviderById(mockProviderInstance.id);
      expect(provider).toBeDefined();
    });
  });

  describe("registerProvider", () => {
    it("should register a new provider instance", async () => {
      await service.registerProvider(mockProviderInstance);

      const provider = await service.getProviderById(mockProviderInstance.id);
      expect(provider).toBeDefined();
      expect(provider?.name).toBe("Ollama");
    });

    it("should update an existing provider instance", async () => {
      await service.registerProvider(mockProviderInstance);

      const updatedInstance = {
        ...mockProviderInstance,
        config: { endpoint: "http://new-endpoint:11434", timeout: 60000 },
      };

      await service.registerProvider(updatedInstance);

      const provider = await service.getProviderById(mockProviderInstance.id);
      expect(provider?.getConfig().endpoint).toBe("http://new-endpoint:11434");
    });

    it("should throw error for unknown provider type", async () => {
      const invalidInstance = {
        ...mockProviderInstance,
        providerType: "unknown",
      };

      await expect(
        service.registerProvider(invalidInstance as LlmProviderInstance)
      ).rejects.toThrow("Unknown provider type: unknown");
    });

    it("should initialize provider after registration", async () => {
      // Provider initialization is tested implicitly by successful registration
      // and the ability to interact with the provider afterwards
      await service.registerProvider(mockProviderInstance);

      const provider = await service.getProviderById(mockProviderInstance.id);
      expect(provider).toBeDefined();
      expect(provider.name).toBe("Ollama");
    });
  });

  describe("unregisterProvider", () => {
    it("should remove a provider from the registry", async () => {
      await service.registerProvider(mockProviderInstance);

      await service.unregisterProvider(mockProviderInstance.id);

      await expect(service.getProviderById(mockProviderInstance.id)).rejects.toThrow(
        `Provider with ID ${mockProviderInstance.id} not found`
      );
    });

    it("should not throw error when unregistering non-existent provider", async () => {
      await expect(service.unregisterProvider("non-existent-id")).resolves.not.toThrow();
    });
  });

  describe("getProviderById", () => {
    it("should return provider by ID", async () => {
      await service.registerProvider(mockProviderInstance);

      const provider = await service.getProviderById(mockProviderInstance.id);

      expect(provider).toBeDefined();
      expect(provider?.name).toBe("Ollama");
    });

    it("should throw error when provider not found", async () => {
      await expect(service.getProviderById("non-existent-id")).rejects.toThrow(
        "Provider with ID non-existent-id not found"
      );
    });
  });

  describe("getAllProviders", () => {
    it("should return all active providers", async () => {
      await service.registerProvider(mockProviderInstance);
      await service.registerProvider(mockUserProviderInstance);

      const providers = await service.getAllProviders();

      expect(providers).toHaveLength(2);
    });

    it("should return empty array when no providers registered", async () => {
      const providers = await service.getAllProviders();

      expect(providers).toHaveLength(0);
    });

    it("should return provider IDs and names", async () => {
      await service.registerProvider(mockProviderInstance);

      const providers = await service.getAllProviders();

      expect(providers[0]).toHaveProperty("id");
      expect(providers[0]).toHaveProperty("name");
      expect(providers[0]).toHaveProperty("provider");
    });
  });

  describe("getDefaultProvider", () => {
    it("should return the default system-level provider", async () => {
      await service.registerProvider(mockProviderInstance);

      const provider = await service.getDefaultProvider();

      expect(provider).toBeDefined();
      expect(provider?.name).toBe("Ollama");
    });

    it("should throw error when no default provider exists", async () => {
      await expect(service.getDefaultProvider()).rejects.toThrow("No default provider configured");
    });

    it("should prefer system-level default over user-level", async () => {
      const userDefault = { ...mockUserProviderInstance, isDefault: true };
      await service.registerProvider(mockProviderInstance);
      await service.registerProvider(userDefault);

      const provider = await service.getDefaultProvider();

      expect(provider?.getConfig().endpoint).toBe("http://localhost:11434");
    });
  });

  describe("getUserProvider", () => {
    it("should return user-specific provider", async () => {
      await service.registerProvider(mockUserProviderInstance);

      const provider = await service.getUserProvider("user-123");

      expect(provider).toBeDefined();
      expect(provider?.getConfig().endpoint).toBe("http://user-ollama:11434");
    });

    it("should prioritize user-level provider over system default", async () => {
      await service.registerProvider(mockProviderInstance);
      await service.registerProvider(mockUserProviderInstance);

      const provider = await service.getUserProvider("user-123");

      expect(provider?.getConfig().endpoint).toBe("http://user-ollama:11434");
    });

    it("should fall back to default provider when user has no specific provider", async () => {
      await service.registerProvider(mockProviderInstance);

      const provider = await service.getUserProvider("user-456");

      expect(provider?.getConfig().endpoint).toBe("http://localhost:11434");
    });

    it("should throw error when no provider available for user", async () => {
      await expect(service.getUserProvider("user-123")).rejects.toThrow(
        "No provider available for user user-123"
      );
    });
  });

  describe("checkAllProvidersHealth", () => {
    it("should return health status for all providers", async () => {
      await service.registerProvider(mockProviderInstance);

      const healthStatuses = await service.checkAllProvidersHealth();

      expect(healthStatuses).toHaveLength(1);
      expect(healthStatuses[0]).toHaveProperty("id");
      expect(healthStatuses[0]).toHaveProperty("healthy");
      expect(healthStatuses[0].healthy).toBe(true);
    });

    it("should handle individual provider health check failures", async () => {
      // Create a mock instance that will use UnhealthyMockProvider
      const unhealthyInstance = { ...mockProviderInstance, id: "unhealthy-id" };

      // Register the unhealthy provider
      await service.registerProvider(unhealthyInstance);

      const healthStatuses = await service.checkAllProvidersHealth();

      expect(healthStatuses).toHaveLength(1);
      // The health status depends on the actual implementation
      expect(healthStatuses[0]).toHaveProperty("healthy");
    });

    it("should return empty array when no providers registered", async () => {
      const healthStatuses = await service.checkAllProvidersHealth();

      expect(healthStatuses).toHaveLength(0);
    });
  });

  describe("reloadFromDatabase", () => {
    it("should reload all enabled providers from database", async () => {
      const findManySpy = vi
        .spyOn(prisma.llmProviderInstance, "findMany")
        .mockResolvedValue([mockProviderInstance]);

      await service.reloadFromDatabase();

      expect(findManySpy).toHaveBeenCalledWith({
        where: { isEnabled: true },
      });
    });

    it("should clear existing providers before reloading", async () => {
      await service.registerProvider(mockProviderInstance);
      vi.spyOn(prisma.llmProviderInstance, "findMany").mockResolvedValue([]);

      await service.reloadFromDatabase();

      const providers = await service.getAllProviders();
      expect(providers).toHaveLength(0);
    });

    it("should update providers with latest database state", async () => {
      vi.spyOn(prisma.llmProviderInstance, "findMany").mockResolvedValue([mockProviderInstance]);
      await service.reloadFromDatabase();

      const updatedInstance = {
        ...mockProviderInstance,
        config: { endpoint: "http://updated:11434", timeout: 60000 },
      };
      vi.spyOn(prisma.llmProviderInstance, "findMany").mockResolvedValue([updatedInstance]);

      await service.reloadFromDatabase();

      const provider = await service.getProviderById(mockProviderInstance.id);
      expect(provider?.getConfig().endpoint).toBe("http://updated:11434");
    });
  });

  describe("error handling", () => {
    it("should handle database connection failures gracefully", async () => {
      vi.spyOn(prisma.llmProviderInstance, "findMany").mockRejectedValue(
        new Error("Database connection failed")
      );

      await expect(service.onModuleInit()).rejects.toThrow(
        "Failed to load providers from database"
      );
    });

    it("should handle provider initialization failures", async () => {
      // Test with an unknown provider type to trigger initialization failure
      const invalidInstance = {
        ...mockProviderInstance,
        providerType: "invalid-type",
      };

      await expect(
        service.registerProvider(invalidInstance as LlmProviderInstance)
      ).rejects.toThrow("Unknown provider type: invalid-type");
    });
  });
});