Epicure/apps/web/lib/ai/__tests__/resolve-user-key.test.ts

import { describe, it, expect, vi, beforeEach } from "vitest";
import { encrypt } from "../../encrypt";

const mockUserAiKeysFindMany = vi.fn();
const mockUserModelPrefsFindFirst = vi.fn();
const mockSiteSettingFindFirst = vi.fn();

vi.mock("@epicure/db", () => ({
  db: {
    query: {
      userAiKeys: { findMany: mockUserAiKeysFindMany },
      userModelPrefs: { findFirst: mockUserModelPrefsFindFirst },
      siteSettings: { findFirst: mockSiteSettingFindFirst },
    },
  },
  userAiKeys: {},
  userModelPrefs: {},
  siteSettings: {},
  eq: vi.fn((a, b) => ({ a, b })),
  and: vi.fn((...args) => args),
}));

const { getDefaultProviderWithKey, getModelConfigForUseCase, withUserKey, ByokDecryptError } = await import("../resolve-user-key");

beforeEach(() => {
  vi.clearAllMocks();
  delete process.env["OPENAI_API_KEY"];
  delete process.env["ANTHROPIC_API_KEY"];
  delete process.env["OPENROUTER_API_KEY"];
  // Make getSiteSetting return null by default
  mockSiteSettingFindFirst.mockResolvedValue(null);
});

describe("getDefaultProviderWithKey", () => {
  it("returns empty config when no keys at all", async () => {
    mockUserAiKeysFindMany.mockResolvedValue([]);
    const config = await getDefaultProviderWithKey("user1");
    expect(config).toEqual({});
  });

  it("returns openrouter key first (priority order)", async () => {
    const encOpenrouter = encrypt("or-key");
    const encOpenai = encrypt("sk-test");
    mockUserAiKeysFindMany.mockResolvedValue([
      { provider: "openai", encryptedKey: encOpenai },
      { provider: "openrouter", encryptedKey: encOpenrouter },
    ]);

    const config = await getDefaultProviderWithKey("user1");
    expect(config.provider).toBe("openrouter");
    expect(config.apiKey).toBe("or-key");
  });

  it("falls back to openai when no openrouter", async () => {
    const encOpenai = encrypt("sk-openai");
    mockUserAiKeysFindMany.mockResolvedValue([
      { provider: "openai", encryptedKey: encOpenai },
    ]);

    const config = await getDefaultProviderWithKey("user1");
    expect(config.provider).toBe("openai");
    expect(config.apiKey).toBe("sk-openai");
  });

  it("uses site settings when no BYOK key exists", async () => {
    mockUserAiKeysFindMany.mockResolvedValue([]);
    // Mock getSiteSetting via siteSettings.findFirst to return encrypted key for OPENAI
    const encKey = encrypt("sk-from-site-settings");
    mockSiteSettingFindFirst.mockResolvedValueOnce(null) // openrouter
      .mockResolvedValueOnce({ value: encKey, isSecret: true }); // openai

    const config = await getDefaultProviderWithKey("user1");
    expect(config.provider).toBe("openai");
    expect(config.apiKey).toBe("sk-from-site-settings");
  });

  it("throws ByokDecryptError instead of silently falling back on a corrupted BYOK key", async () => {
    mockUserAiKeysFindMany.mockResolvedValue([
      { provider: "openrouter", encryptedKey: "CORRUPT:NOT:VALID" },
      { provider: "openai", encryptedKey: encrypt("sk-valid") },
    ]);

    await expect(getDefaultProviderWithKey("user1")).rejects.toThrow(ByokDecryptError);
  });
});

describe("withUserKey", () => {
  it("injects BYOK key when user has one for this provider", async () => {
    const encKey = encrypt("sk-user-key");
    vi.mocked(mockUserAiKeysFindMany); // just ensure mock is ready
    // withUserKey uses findFirst via userAiKeys
    const mockFindFirst = vi.fn().mockResolvedValue({ encryptedKey: encKey });
    ((await import("@epicure/db")).db.query.userAiKeys as unknown as { findFirst: typeof mockFindFirst }).findFirst = mockFindFirst;

    const config = await withUserKey("user1", { provider: "openai" });
    expect(config.apiKey).toBe("sk-user-key");
  });

  it("returns config unchanged when no user key for provider", async () => {
    const mockFindFirst = vi.fn().mockResolvedValue(null);
    ((await import("@epicure/db")).db.query.userAiKeys as unknown as { findFirst: typeof mockFindFirst }).findFirst = mockFindFirst;

    const config = await withUserKey("user1", { provider: "anthropic" });
    expect(config.apiKey).toBeUndefined();
  });
});

describe("getModelConfigForUseCase", () => {
  it("uses user model prefs when set", async () => {
    mockUserModelPrefsFindFirst.mockResolvedValue({
      textProvider: "anthropic",
      textModel: "claude-sonnet-4-6",
      visionProvider: null,
      visionModel: null,
      mealPlanProvider: null,
      mealPlanModel: null,
    });
    const mockFindFirst = vi.fn().mockResolvedValue(null); // no BYOK key
    ((await import("@epicure/db")).db.query.userAiKeys as unknown as { findFirst: typeof mockFindFirst }).findFirst = mockFindFirst;

    const config = await getModelConfigForUseCase("user1", "text");
    expect(config.provider).toBe("anthropic");
    expect(config.model).toBe("claude-sonnet-4-6");
  });

  it("falls back to getDefaultProviderWithKey when no prefs", async () => {
    mockUserModelPrefsFindFirst.mockResolvedValue(null);
    mockUserAiKeysFindMany.mockResolvedValue([
      { provider: "openai", encryptedKey: encrypt("sk-default") },
    ]);

    const config = await getModelConfigForUseCase("user1", "vision");
    expect(config.provider).toBe("openai");
    expect(config.apiKey).toBe("sk-default");
  });
});