Alfred/tests/server/ai/registry.test.ts

import { describe, expect, mock, test } from "bun:test";

import { createNoopLogger } from "../../../src/server/logger";
import { createMigratedTestDatabase } from "../../helpers";

void mock.module("ai", () => ({
  createProviderRegistry: (providers: Record<string, { languageModel: (modelId: string) => unknown }>) => ({
    languageModel: (id: string) => {
      const [providerId, modelId] = id.split(":");
      const provider = providers[providerId ?? ""];
      if (!provider || !modelId) throw new Error(`No such provider: ${id}`);
      return provider.languageModel(modelId);
    },
  }),
  generateText: () => Promise.resolve({ text: "Hi" }),
}));

async function withProviderServer(
  modelsResponse: Response,
  callback: (baseUrl: string) => Promise<void>,
): Promise<void> {
  const server = Bun.serve({
    fetch(request) {
      if (request.method === "HEAD") return new Response(null, { status: 200 });
      return modelsResponse;
    },
    port: 0,
  });
  try {
    await callback(`http://127.0.0.1:${server.port}/v1`);
  } finally {
    await server.stop(true);
  }
}

describe("AI registry", () => {
  test("testProviderConnection reports unreachable Base URL", async () => {
    const { testProviderConnection } = await import("../../../src/server/ai/registry");

    const result = await testProviderConnection(
      {
        apiKey: "bad-key",
        baseUrl: "http://127.0.0.1:1",
        name: "Bad",
        type: "openai-compatible",
      },
      createNoopLogger(),
    );

    expect(result.ok).toBe(false);
    expect(result.message).toContain("Base URL 不可达");
  });

  test("testProviderConnection rejects invalid config", async () => {
    await withProviderServer(new Response(null, { status: 401 }), async (baseUrl) => {
      const { testProviderConnection } = await import("../../../src/server/ai/registry");

      const result = await testProviderConnection(
        {
          apiKey: "bad-key",
          baseUrl,
          name: "Bad",
          type: "openai-compatible",
        },
        createNoopLogger(),
      );

      expect(result.ok).toBe(false);
      expect(result.message).toContain("API Key 无效");
      expect(typeof result.message).toBe("string");
    });
  });

  test("testProviderConnection return shape is correct", async () => {
    await withProviderServer(Response.json({ data: [{ id: "gpt-4o" }] }), async (baseUrl) => {
      const { testProviderConnection } = await import("../../../src/server/ai/registry");

      const result = await testProviderConnection(
        {
          apiKey: "sk-test",
          baseUrl,
          name: "Test",
          type: "openai",
        },
        createNoopLogger(),
      );

      expect(result.ok).toBe(true);
      expect(result.message).toContain("/models 返回 1 个模型");
    });
  });

  test("testProviderConnection treats unsupported /models as non-blocking", async () => {
    await withProviderServer(new Response(null, { status: 404 }), async (baseUrl) => {
      const { testProviderConnection } = await import("../../../src/server/ai/registry");

      const result = await testProviderConnection(
        {
          apiKey: "sk-test",
          baseUrl,
          name: "Test",
          type: "openai",
        },
        createNoopLogger(),
      );

      expect(result.ok).toBe(true);
      expect(result.message).toContain("可能不支持 /models");
    });
  });

  test("buildProviderRegistry 从 DB 构建包含所有供应商的注册表", async () => {
    const handle = createMigratedTestDatabase("registry-build-test");
    const now = new Date().toISOString();

    handle.db
      .prepare(
        "INSERT INTO providers (id, name, type, base_url, api_key, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
      )
      .run("pv1", "OpenAI", "openai", "https://api.openai.com/v1", "sk-test", now, now);
    handle.db
      .prepare(
        "INSERT INTO providers (id, name, type, base_url, api_key, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
      )
      .run("pv2", "Anthropic", "anthropic", "https://api.anthropic.com", "sk-off", now, now);

    const { buildProviderRegistry } = await import("../../../src/server/ai/registry");
    const registry = buildProviderRegistry(handle.db);

    expect(() => registry.languageModel("pv1:gpt-4o")).not.toThrow();
    expect(() => registry.languageModel("pv2:claude-3")).not.toThrow();

    handle.cleanup();
  });

  test("buildProviderRegistry 无供应商时返回空注册表", async () => {
    const handle = createMigratedTestDatabase("registry-empty-test");

    const { buildProviderRegistry } = await import("../../../src/server/ai/registry");
    const registry = buildProviderRegistry(handle.db);

    expect(typeof registry.languageModel).toBe("function");

    handle.cleanup();
  });

  test("testModelConnection 成功返回 ok:true", async () => {
    const { testModelConnection } = await import("../../../src/server/ai/registry");

    const result = await testModelConnection(
      {
        apiKey: "sk-test",
        baseUrl: "https://api.openai.com/v1",
        modelId: "gpt-4o",
        name: "Test",
        type: "openai",
      },
      createNoopLogger(),
    );

    expect(result.ok).toBe(true);
    expect(result.message).toContain("模型连接成功");
  });
});