Alfred/tests/server/db/models.test.ts

import type Database from "bun:sqlite";

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

import {
  createModel,
  deleteModel,
  disableModel,
  enableModel,
  getModel,
  getModelsByProviderId,
  listModels,
  updateModel,
} from "../../../src/server/db/models";
import { createProvider } from "../../../src/server/db/providers";
import { createMigratedTestDatabase } from "../../helpers";

function seedProvider(db: Database, name = "TestProvider"): string {
  const result = createProvider(db, { apiKey: "sk-test", baseUrl: "https://api.test.com/v1", name, type: "openai" });
  return (result as { provider: { id: string } }).provider.id;
}

function withDb(callback: (db: Database) => void): void {
  const handle = createMigratedTestDatabase("models-test");
  try {
    callback(handle.db);
    handle.close();
  } finally {
    handle.cleanup();
  }
}

describe("模型数据访问层", () => {
  test("创建模型", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const result = createModel(db, {
        capabilities: ["text", "reasoning"],
        modelId: "gpt-4o",
        name: "GPT-4o",
        providerId,
      });
      expect("error" in result).toBe(false);
      const model = (
        result as {
          model: { capabilities: string[]; enabled: boolean; modelId: string; name: string; providerId: string };
        }
      ).model;
      expect(model.name).toBe("GPT-4o");
      expect(model.modelId).toBe("gpt-4o");
      expect(model.providerId).toBe(providerId);
      expect(model.capabilities).toEqual(["text", "reasoning"]);
      expect(model.enabled).toBe(true);
    });
  });

  test("供应商不存在时创建失败", () => {
    withDb((db) => {
      const result = createModel(db, {
        capabilities: ["text"],
        modelId: "test",
        name: "Test",
        providerId: "nonexistent",
      });
      expect("error" in result).toBe(true);
      expect((result as unknown as { status: number }).status).toBe(400);
    });
  });

  test("同一供应商下模型 ID 唯一", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      createModel(db, { capabilities: ["text"], modelId: "gpt-4o", name: "Model1", providerId });
      const result = createModel(db, { capabilities: ["text"], modelId: "gpt-4o", name: "Model2", providerId });
      expect("error" in result).toBe(true);
      expect((result as unknown as { error: string }).error).toContain("已存在");
    });
  });

  test("不同供应商下模型 ID 可重复", () => {
    withDb((db) => {
      const p1 = seedProvider(db, "P1");
      const p2 = seedProvider(db, "P2");
      const r1 = createModel(db, { capabilities: ["text"], modelId: "same-id", name: "M1", providerId: p1 });
      const r2 = createModel(db, { capabilities: ["text"], modelId: "same-id", name: "M2", providerId: p2 });
      expect("error" in r1).toBe(false);
      expect("error" in r2).toBe(false);
    });
  });

  test("能力标签为空时创建失败", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const result = createModel(db, { capabilities: [], modelId: "test", name: "Test", providerId });
      expect("error" in result).toBe(true);
      expect((result as unknown as { error: string }).error).toContain("能力标签");
    });
  });

  test("列表查询（分页、关键字、供应商过滤）", () => {
    withDb((db) => {
      const p1 = seedProvider(db, "P1");
      const p2 = seedProvider(db, "P2");
      createModel(db, { capabilities: ["text"], modelId: "m1", name: "Alpha", providerId: p1 });
      createModel(db, { capabilities: ["text"], modelId: "m2", name: "Beta", providerId: p1 });
      createModel(db, { capabilities: ["text"], modelId: "m3", name: "Gamma", providerId: p2 });

      const all = listModels(db, { page: 1, pageSize: 20 });
      expect(all.total).toBe(3);

      const filtered = listModels(db, { page: 1, pageSize: 20, providerId: p1 });
      expect(filtered.total).toBe(2);

      const searched = listModels(db, { keyword: "Alpha", page: 1, pageSize: 20 });
      expect(searched.total).toBe(1);
    });
  });

  test("获取模型详情", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const created = createModel(db, { capabilities: ["text"], modelId: "gpt-4o", name: "GPT-4o", providerId });
      const id = (created as { model: { id: string } }).model.id;

      const result = getModel(db, id);
      expect("error" in result).toBe(false);
      expect((result as { model: { name: string } }).model.name).toBe("GPT-4o");
    });
  });

  test("获取不存在的模型返回 404", () => {
    withDb((db) => {
      const result = getModel(db, "nonexistent");
      expect("error" in result).toBe(true);
      expect((result as unknown as { status: number }).status).toBe(404);
    });
  });

  test("更新模型", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const created = createModel(db, { capabilities: ["text"], modelId: "gpt-4o", name: "原名", providerId });
      const id = (created as { model: { id: string } }).model.id;

      const result = updateModel(db, id, { capabilities: ["text", "reasoning"], name: "新名" });
      expect("error" in result).toBe(false);
      const updated = (result as { model: { capabilities: string[]; name: string } }).model;
      expect(updated.name).toBe("新名");
      expect(updated.capabilities).toEqual(["text", "reasoning"]);
    });
  });

  test("启用/禁用模型", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const created = createModel(db, { capabilities: ["text"], modelId: "gpt-4o", name: "测试", providerId });
      const id = (created as { model: { id: string } }).model.id;

      const disabled = disableModel(db, id);
      expect("error" in disabled).toBe(false);
      expect((disabled as { model: { enabled: boolean } }).model.enabled).toBe(false);

      const enabled = enableModel(db, id);
      expect("error" in enabled).toBe(false);
      expect((enabled as { model: { enabled: boolean } }).model.enabled).toBe(true);
    });
  });

  test("重复禁用失败", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const created = createModel(db, { capabilities: ["text"], modelId: "gpt-4o", name: "测试", providerId });
      const id = (created as { model: { id: string } }).model.id;
      disableModel(db, id);

      const result = disableModel(db, id);
      expect("error" in result).toBe(true);
      expect((result as unknown as { status: number }).status).toBe(409);
    });
  });

  test("删除模型", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const created = createModel(db, { capabilities: ["text"], modelId: "gpt-4o", name: "删除测试", providerId });
      const id = (created as { model: { id: string } }).model.id;

      const result = deleteModel(db, id);
      expect("error" in result).toBe(false);

      const after = getModel(db, id);
      expect("error" in after).toBe(true);
    });
  });

  test("getModelsByProviderId 返回正确数量", () => {
    withDb((db) => {
      const p1 = seedProvider(db, "P1");
      const p2 = seedProvider(db, "P2");
      createModel(db, { capabilities: ["text"], modelId: "m1", name: "M1", providerId: p1 });
      createModel(db, { capabilities: ["text"], modelId: "m2", name: "M2", providerId: p1 });
      createModel(db, { capabilities: ["text"], modelId: "m3", name: "M3", providerId: p2 });

      expect(getModelsByProviderId(db, p1)).toBe(2);
      expect(getModelsByProviderId(db, p2)).toBe(1);
    });
  });

  test("可选字段 contextLength 和 maxOutputTokens", () => {
    withDb((db) => {
      const providerId = seedProvider(db);
      const result = createModel(db, {
        capabilities: ["text"],
        contextLength: 128000,
        maxOutputTokens: 4096,
        modelId: "gpt-4o",
        name: "GPT-4o",
        providerId,
      });
      expect("error" in result).toBe(false);
      const model = (result as { model: { contextLength: null | number; maxOutputTokens: null | number } }).model;
      expect(model.contextLength).toBe(128000);
      expect(model.maxOutputTokens).toBe(4096);
    });
  });
});