next-ai-draw-io/tests/unit/use-model-config.test.tsx

import { act, cleanup, renderHook, waitFor } from "@testing-library/react"
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import { getSelectedAIConfig, useModelConfig } from "@/hooks/use-model-config"
import type { FlattenedServerModel } from "@/lib/server-model-config"
import { STORAGE_KEYS } from "@/lib/storage"
import type { MultiModelConfig } from "@/lib/types/model-config"

const SERVER_MODELS: FlattenedServerModel[] = [
    {
        id: "server:openai-main:gpt-4o-mini",
        modelId: "gpt-4o-mini",
        provider: "openai",
        providerLabel: "OpenAI Main",
        isDefault: false,
    },
    {
        id: "server:openai-main:gpt-4o",
        modelId: "gpt-4o",
        provider: "openai",
        providerLabel: "OpenAI Main",
        isDefault: true,
    },
]

const USER_CONFIG: MultiModelConfig = {
    version: 1,
    providers: [
        {
            id: "p1",
            provider: "openai",
            apiKey: "sk-test",
            models: [{ id: "m1", modelId: "gpt-4o" }],
        },
    ],
}

function storeConfig(config: MultiModelConfig) {
    localStorage.setItem(STORAGE_KEYS.modelConfigs, JSON.stringify(config))
}

async function renderLoaded() {
    const hook = renderHook(() => useModelConfig())
    await waitFor(() => expect(hook.result.current.isLoaded).toBe(true))
    return hook
}

beforeEach(() => {
    localStorage.clear()
    vi.stubGlobal(
        "fetch",
        vi.fn(async () => ({
            ok: true,
            json: async () => ({ models: SERVER_MODELS }),
        })),
    )
})

afterEach(() => {
    cleanup()
    vi.unstubAllGlobals()
})

describe("useModelConfig server model selection", () => {
    it("replaces a saved server model that no longer exists", async () => {
        storeConfig({
            ...USER_CONFIG,
            selectedModelId: "server:openai-production:gpt-4o",
        })
        const { result } = await renderLoaded()
        expect(result.current.selectedModelId).toBe("server:openai-main:gpt-4o")
    })

    it("keeps a server model saved under its old id", async () => {
        // Before non-ASCII provider names got their own slug, "主力 OpenAI"
        // became "openai"
        const renamed: FlattenedServerModel = {
            id: "server:4e3b-529b-openai:gpt-4o-mini",
            modelId: "gpt-4o-mini",
            provider: "openai",
            providerLabel: "主力 OpenAI",
            isDefault: false,
        }
        vi.stubGlobal(
            "fetch",
            vi.fn(async () => ({
                ok: true,
                json: async () => ({ models: [...SERVER_MODELS, renamed] }),
            })),
        )
        storeConfig({
            ...USER_CONFIG,
            selectedModelId: "server:openai:gpt-4o-mini",
        })
        const { result } = await renderLoaded()
        expect(result.current.selectedModelId).toBe(renamed.id)
    })

    it("keeps a saved server model that still exists", async () => {
        storeConfig({
            ...USER_CONFIG,
            selectedModelId: "server:openai-main:gpt-4o-mini",
        })
        const { result } = await renderLoaded()
        expect(result.current.selectedModelId).toBe(
            "server:openai-main:gpt-4o-mini",
        )
    })

    it("skips a saved provider this version does not know", async () => {
        // Saved by another version, or edited by hand: it used to crash the
        // whole page on load
        storeConfig({
            ...USER_CONFIG,
            providers: [
                ...USER_CONFIG.providers,
                {
                    id: "p9",
                    provider: "not-a-provider" as any,
                    apiKey: "k",
                    models: [{ id: "m9", modelId: "x" }],
                },
            ],
        })
        const { result } = await renderLoaded()
        expect(result.current.config.providers.map((p) => p.id)).toEqual(["p1"])
        expect(result.current.models.map((m) => m.id)).toContain("m1")
    })

    it("keeps an unknown provider and its key in storage", async () => {
        // The version that saved it may be opened again (an older desktop
        // build, another tab): the provider must still be there
        storeConfig({
            ...USER_CONFIG,
            providers: [
                ...USER_CONFIG.providers,
                {
                    id: "p9",
                    provider: "not-a-provider" as any,
                    apiKey: "k9",
                    models: [{ id: "m9", modelId: "x" }],
                },
            ],
            selectedModelId: "m1",
        })
        const { result } = await renderLoaded()
        act(() => result.current.setSelectedModelId(undefined))
        await waitFor(() => {
            const stored = JSON.parse(
                localStorage.getItem(STORAGE_KEYS.modelConfigs) ?? "{}",
            )
            expect(stored.selectedModelId).toBeUndefined()
            expect(stored.providers.map((p: { id: string }) => p.id)).toEqual([
                "p1",
                "p9",
            ])
            expect(stored.providers[1].apiKey).toBe("k9")
        })
        // Sending reads the stored config too, and must not trip over it
        act(() => result.current.setSelectedModelId("m1"))
        expect(getSelectedAIConfig()).toMatchObject({
            aiProvider: "openai",
            aiModel: "gpt-4o",
        })
    })

    it("keeps a selected user model", async () => {
        storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
        const { result } = await renderLoaded()
        expect(result.current.selectedModelId).toBe("m1")
    })

    it("falls back to the default server model when the selected model is deleted", async () => {
        storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
        const { result } = await renderLoaded()
        act(() => result.current.deleteModel("p1", "m1"))
        expect(result.current.selectedModelId).toBe("server:openai-main:gpt-4o")
    })

    it("falls back to the default server model when the selected provider is deleted", async () => {
        storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
        const { result } = await renderLoaded()
        act(() => result.current.deleteProvider("p1"))
        expect(result.current.selectedModelId).toBe("server:openai-main:gpt-4o")
    })
})

describe("useModelConfig in the desktop app", () => {
    it("reloads the server models after a preset switch restarts the server", async () => {
        // The new preset offers other models; the saved one is gone
        let restarted: (() => void) | undefined
        ;(window as any).electronAPI = {
            onServerRestarted: (callback: () => void) => {
                restarted = callback
                return () => {
                    restarted = undefined
                }
            },
        }
        try {
            storeConfig({
                ...USER_CONFIG,
                selectedModelId: "server:openai-main:gpt-4o-mini",
            })
            const { result } = await renderLoaded()
            await waitFor(() => expect(restarted).toBeDefined())
            const nextModels: FlattenedServerModel[] = [
                {
                    id: "server:claude:claude-sonnet-5-5",
                    modelId: "claude-sonnet-5-5",
                    provider: "anthropic",
                    providerLabel: "Claude",
                    isDefault: true,
                },
            ]
            vi.stubGlobal(
                "fetch",
                vi.fn(async () => ({
                    ok: true,
                    json: async () => ({ models: nextModels }),
                })),
            )
            act(() => restarted?.())
            await waitFor(() =>
                expect(result.current.selectedModelId).toBe(
                    "server:claude:claude-sonnet-5-5",
                ),
            )
        } finally {
            delete (window as any).electronAPI
        }
    })
})

describe("useModelConfig across tabs", () => {
    it("reloads the config when another tab saves it", async () => {
        storeConfig({ ...USER_CONFIG, selectedModelId: "m1" })
        const { result } = await renderLoaded()

        const fromOtherTab: MultiModelConfig = {
            ...USER_CONFIG,
            providers: [
                ...USER_CONFIG.providers,
                {
                    id: "p2",
                    provider: "anthropic",
                    apiKey: "sk-ant",
                    models: [{ id: "m2", modelId: "claude-sonnet-4-5" }],
                },
            ],
            selectedModelId: "m2",
        }
        act(() => {
            storeConfig(fromOtherTab)
            window.dispatchEvent(
                new StorageEvent("storage", { key: STORAGE_KEYS.modelConfigs }),
            )
        })

        expect(result.current.selectedModelId).toBe("m2")
        expect(result.current.config.providers).toHaveLength(2)
    })
})

describe("useModelConfig migration", () => {
    const OLD = {
        provider: "next-ai-draw-io-ai-provider",
        baseUrl: "next-ai-draw-io-ai-base-url",
        apiKey: "next-ai-draw-io-ai-api-key",
        model: "next-ai-draw-io-ai-model",
    }

    it("moves an old keyless Ollama setup into the new format", async () => {
        localStorage.setItem(OLD.provider, "ollama")
        localStorage.setItem(OLD.baseUrl, "http://localhost:11434/api")
        localStorage.setItem(OLD.apiKey, "")
        localStorage.setItem(OLD.model, "llama3.2")
        const { result } = await renderLoaded()
        const [provider] = result.current.config.providers
        expect(provider).toMatchObject({
            provider: "ollama",
            apiKey: "",
            baseUrl: "http://localhost:11434/api",
            models: [{ modelId: "llama3.2" }],
        })
        expect(result.current.selectedModelId).toBe(provider.models[0].id)
        expect(localStorage.getItem(OLD.model)).toBeNull()
    })

    it("keeps an old Ollama without a URL on the server's Ollama", async () => {
        localStorage.setItem(OLD.provider, "ollama")
        localStorage.setItem(OLD.apiKey, "")
        localStorage.setItem(OLD.model, "llama3.2")
        const { result } = await renderLoaded()
        // No cloud address: the request takes OLLAMA_BASE_URL or this machine
        expect(result.current.config.providers[0].baseUrl).toBe("")
    })

    it("still needs a key for other providers", async () => {
        localStorage.setItem(OLD.provider, "openai")
        localStorage.setItem(OLD.model, "gpt-4o")
        const { result } = await renderLoaded()
        expect(result.current.config.providers).toEqual([])
    })
})