Fix model test routing for image providers

tests/unit/hf-model-routing.test.js

@@ -0,0 +1,12 @@
+import { describe, it, expect } from "vitest";
+import { parseModel } from "../../open-sse/services/model.js";
+
+describe("HuggingFace model alias parsing", () => {
+  it("resolves hf alias to huggingface provider", () => {
+    expect(parseModel("hf/black-forest-labs/FLUX.1-schnell")).toMatchObject({
+      provider: "huggingface",
+      model: "black-forest-labs/FLUX.1-schnell",
+      providerAlias: "hf",
+    });
+  });
+});

tests/unit/model-test-routing.test.js

@@ -0,0 +1,74 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  getApiKeys: vi.fn(),
+  getConsistentMachineId: vi.fn(),
+}));
+
+vi.mock("@/lib/localDb", () => ({
+  getApiKeys: mocks.getApiKeys,
+}));
+
+vi.mock("@/shared/utils/machineId", () => ({
+  getConsistentMachineId: mocks.getConsistentMachineId,
+}));
+
+vi.mock("next/server", () => ({
+  NextResponse: {
+    json(body, init = {}) {
+      return new Response(JSON.stringify(body), {
+        status: init.status || 200,
+        headers: { "Content-Type": "application/json" },
+      });
+    },
+  },
+}));
+
+const originalFetch = global.fetch;
+
+describe("model test route kind routing", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mocks.getApiKeys.mockResolvedValue([{ key: "sk-internal", isActive: true }]);
+    mocks.getConsistentMachineId.mockResolvedValue("cli-token");
+    global.fetch = vi.fn().mockResolvedValue(new Response(JSON.stringify({
+      created: 1,
+      data: [{ b64_json: "abc" }],
+    }), {
+      status: 200,
+      headers: { "Content-Type": "application/json" },
+    }));
+  });
+
+  afterEach(() => {
+    global.fetch = originalFetch;
+  });
+
+  it("routes image model tests to /api/v1/images/generations", async () => {
+    const { POST } = await import("../../src/app/api/models/test/route.js");
+
+    const req = new Request("http://localhost/api/models/test", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({
+        model: "hf/black-forest-labs/FLUX.1-schnell",
+        kind: "image",
+      }),
+    });
+
+    const res = await POST(req);
+    const body = await res.json();
+
+    expect(body.ok).toBe(true);
+    expect(global.fetch).toHaveBeenCalledWith(
+      expect.stringContaining("/api/v1/images/generations"),
+      expect.objectContaining({
+        method: "POST",
+        body: JSON.stringify({
+          model: "hf/black-forest-labs/FLUX.1-schnell",
+          prompt: "test",
+        }),
+      })
+    );
+  });
+});

tests/unit/provider-test-models-routing.test.js

@@ -0,0 +1,83 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  getProviderConnectionById: vi.fn(),
+  getApiKeys: vi.fn(),
+  getConsistentMachineId: vi.fn(),
+}));
+
+vi.mock("@/lib/localDb", () => ({
+  getProviderConnectionById: mocks.getProviderConnectionById,
+  getApiKeys: mocks.getApiKeys,
+}));
+
+vi.mock("@/shared/utils/machineId", () => ({
+  getConsistentMachineId: mocks.getConsistentMachineId,
+}));
+
+vi.mock("next/server", () => ({
+  NextResponse: {
+    json(body, init = {}) {
+      return new Response(JSON.stringify(body), {
+        status: init.status || 200,
+        headers: { "Content-Type": "application/json" },
+      });
+    },
+  },
+}));
+
+const originalFetch = global.fetch;
+
+describe("provider test-models route kind routing", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mocks.getProviderConnectionById.mockResolvedValue({
+      id: "conn-hf",
+      provider: "huggingface",
+    });
+    mocks.getApiKeys.mockResolvedValue([{ key: "sk-internal", isActive: true }]);
+    mocks.getConsistentMachineId.mockResolvedValue("cli-token");
+    global.fetch = vi.fn((url) => {
+      if (String(url).includes("/api/v1/images/generations")) {
+        return Promise.resolve(new Response(JSON.stringify({
+          created: 1,
+          data: [{ b64_json: "abc" }],
+        }), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        }));
+      }
+      return Promise.resolve(new Response(JSON.stringify({
+        choices: [{ message: { role: "assistant", content: "ok" } }],
+      }), {
+        status: 200,
+        headers: { "Content-Type": "application/json" },
+      }));
+    });
+  });
+
+  afterEach(() => {
+    global.fetch = originalFetch;
+  });
+
+  it("routes huggingface image models to /api/v1/images/generations", async () => {
+    const { POST } = await import("../../src/app/api/providers/[id]/test-models/route.js");
+
+    const req = new Request("http://localhost/api/providers/conn-hf/test-models", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+    });
+
+    const res = await POST(req, { params: Promise.resolve({ id: "conn-hf" }) });
+    const body = await res.json();
+
+    expect(body.provider).toBe("huggingface");
+    expect(body.results.some((r) => r.modelId === "black-forest-labs/FLUX.1-schnell" && r.ok)).toBe(true);
+    expect(global.fetch).toHaveBeenCalledWith(
+      expect.stringContaining("/api/v1/images/generations"),
+      expect.objectContaining({
+        method: "POST",
+      })
+    );
+  });
+});