popiart-node-canvas/src/app/api/generate/providers/__tests__/newapiwg.test.ts

import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import {
  fetchNewApiWGModels,
  generateWithNewApiWG,
  inferNewApiWGCapabilities,
  normalizeNewApiWGError,
} from "../newapiwg";
import type { GenerationInput } from "@/lib/providers/types";

let capturedBody: Record<string, unknown> | null = null;

function makeInput(overrides: Partial<GenerationInput> = {}): GenerationInput {
  return {
    model: {
      id: "doubao-seedance-2-0-260128",
      name: "Seedance",
      description: null,
      provider: "newapiwg",
      capabilities: ["text-to-video", "image-to-video", "audio-to-video"],
    },
    prompt: "a dancing robot",
    images: [],
    parameters: {},
    ...overrides,
  };
}

describe("NewApiWG model capability inference", () => {
  it("classifies real NewAPI gateway model ids without treating chat models as image generators", () => {
    expect(inferNewApiWGCapabilities({
      id: "apiyi_sora_image",
      owned_by: "custom",
      supported_endpoint_types: ["openai"],
    })).toEqual(["text-to-image"]);

    expect(inferNewApiWGCapabilities({
      id: "gpt-image-2-all",
      owned_by: "openai",
      supported_endpoint_types: ["openai"],
    })).toEqual(["text-to-image", "image-to-image"]);

    expect(inferNewApiWGCapabilities({
      id: "apiyi_nano_banana_pro",
      owned_by: "custom",
      supported_endpoint_types: ["gemini", "openai"],
    })).toEqual(["text-to-image", "image-to-image"]);

    expect(inferNewApiWGCapabilities({
      id: "gemini-3-pro-image-preview",
      owned_by: "custom",
      supported_endpoint_types: ["gemini", "openai"],
    })).toEqual(["text-to-image", "image-to-image"]);

    expect(inferNewApiWGCapabilities({
      id: "doubao-seedance-2-0-260128",
      owned_by: "doubao-video",
      supported_endpoint_types: ["openai"],
    })).toEqual(["text-to-video", "image-to-video"]);

    expect(inferNewApiWGCapabilities({
      id: "speech-2.8-hd",
      owned_by: "minimax",
      supported_endpoint_types: ["openai"],
    })).toEqual(["text-to-audio"]);

    expect(inferNewApiWGCapabilities({
      id: "kimi-k2.6",
      owned_by: "moonshot",
      supported_endpoint_types: ["openai"],
    })).toEqual([]);

    expect(inferNewApiWGCapabilities({
      id: "popiart-custom-model",
      owned_by: "popiart",
      supported_endpoint_types: ["openai"],
    })).toEqual([
      "text-to-image",
      "image-to-image",
      "text-to-video",
      "image-to-video",
      "text-to-3d",
      "image-to-3d",
      "text-to-audio",
      "audio-to-video",
    ]);
  });

  it("omits pure chat models from generation model discovery", async () => {
    vi.stubGlobal("fetch", vi.fn(async () => new Response(JSON.stringify({
      data: [
        { id: "kimi-k2.6", owned_by: "moonshot", supported_endpoint_types: ["openai"] },
        { id: "doubao-seedance-2-0-260128", owned_by: "doubao-video", supported_endpoint_types: ["openai"] },
        { id: "apiyi_sora_image", owned_by: "custom", supported_endpoint_types: ["openai"] },
      ],
    }), { status: 200 })));

    const models = await fetchNewApiWGModels("test-key", "https://newapi.example/v1");

    expect(models.map((model) => model.id)).toEqual([
      "doubao-seedance-2-0-260128",
      "apiyi_sora_image",
    ]);
  });
});

describe("NewApiWG generation payloads", () => {
  beforeEach(() => {
    capturedBody = null;
    vi.stubGlobal("fetch", vi.fn(async (url: string | URL | Request, init?: RequestInit) => {
      const urlStr = typeof url === "string" ? url : url instanceof URL ? url.toString() : url.url;

      if (urlStr === "https://cdn.example.com/seedream.png") {
        return new Response(Buffer.from("abc"), {
          status: 200,
          headers: { "content-type": "image/png" },
        });
      }

      if (urlStr.includes("/v1beta/models/") && urlStr.includes(":generateContent") && init?.method === "POST") {
        capturedBody = JSON.parse(init.body as string);
        return new Response(JSON.stringify({
          candidates: [{
            content: {
              parts: [{
                inlineData: {
                  mimeType: "image/png",
                  data: "gemini123",
                },
              }],
            },
          }],
        }), {
          status: 200,
        });
      }

      if (urlStr.includes("/chat/completions") && init?.method === "POST") {
        capturedBody = JSON.parse(init.body as string);
        return new Response(JSON.stringify({
          choices: [{
            message: {
              content: [{
                type: "image_url",
                image_url: { url: "data:image/png;base64,abc123" },
              }],
            },
          }],
        }), {
          status: 200,
        });
      }

      if (urlStr.includes("/video/generations") && init?.method === "POST") {
        capturedBody = JSON.parse(init.body as string);
        return new Response(JSON.stringify({ video_url: "https://cdn.example.com/video.mp4" }), {
          status: 200,
        });
      }

      if (urlStr.includes("/images/generations") && init?.method === "POST") {
        capturedBody = JSON.parse(init.body as string);
        return new Response(JSON.stringify({
          data: [{ url: "https://cdn.example.com/seedream.png" }],
        }), {
          status: 200,
        });
      }

      return new Response("Not Found", { status: 404 });
    }));
  });

  afterEach(() => {
    vi.unstubAllGlobals();
    vi.useRealTimers();
  });

  it("passes schema-mapped dynamic image and audio inputs to video generation", async () => {
    const result = await generateWithNewApiWG("test-key", "https://newapi.example/v1", makeInput({
      images: ["data:image/png;base64,abc"],
      parameters: { duration: 5 },
      dynamicInputs: {
        first_frame_url: "data:image/png;base64,first",
        reference_audio_urls: ["data:audio/mpeg;base64,audio"],
      },
    }));

    expect(result.success).toBe(true);
    expect(capturedBody).not.toBeNull();
    expect(capturedBody!.prompt).toBe("a dancing robot");
    expect(capturedBody!.images).toEqual(["data:image/png;base64,abc"]);
    expect(capturedBody!.first_frame_url).toBe("data:image/png;base64,first");
    expect(capturedBody!.reference_audio_urls).toEqual(["data:audio/mpeg;base64,audio"]);
    expect(capturedBody!.duration).toBe(5);
  });

  it("routes nano banana image models through Gemini native generation", async () => {
    const result = await generateWithNewApiWG("test-key", "https://newapi.example/v1", makeInput({
      model: {
        id: "apiyi_nano_banana_2",
        name: "apiyi_nano_banana_2",
        description: null,
        provider: "newapiwg",
        capabilities: ["text-to-image", "image-to-image"],
      },
      prompt: "draw a banana",
      images: ["data:image/png;base64,input"],
    }));

    expect(result.success).toBe(true);
    expect(result.outputs?.[0]).toEqual({ type: "image", data: "data:image/png;base64,gemini123" });
    expect(capturedBody!.contents).toEqual([
      {
        parts: [
          { text: "draw a banana" },
          { inlineData: { mimeType: "image/png", data: "input" } },
        ]
      },
    ]);
    expect(capturedBody!.generationConfig).toMatchObject({
      responseModalities: ["IMAGE"],
      imageConfig: { aspectRatio: "1:1", imageSize: "2K" },
    });
  });

  it("accepts Gemini native fileData URLs without image file extensions", async () => {
    vi.stubGlobal("fetch", vi.fn(async (url: string | URL | Request, init?: RequestInit) => {
      capturedBody = init?.body ? JSON.parse(init.body as string) : null;
      return new Response(JSON.stringify({
        candidates: [{
          content: {
            parts: [{
              fileData: {
                mimeType: "image/png",
                fileUri: "https://cdn.example.com/generated/abc123?signature=test",
              },
            }],
          },
        }],
      }), { status: 200 });
    }));

    const result = await generateWithNewApiWG("test-key", "https://newapi.example/v1", makeInput({
      model: {
        id: "apiyi_nano_banana_2",
        name: "apiyi_nano_banana_2",
        description: null,
        provider: "newapiwg",
        capabilities: ["text-to-image", "image-to-image"],
      },
      prompt: "draw a banana",
    }));

    expect(result.success).toBe(true);
    expect(result.outputs?.[0]).toEqual({
      type: "image",
      data: "",
      url: "https://cdn.example.com/generated/abc123?signature=test",
    });
  });

  it("retries transient overload errors from Gemini native image generation", async () => {
    vi.useFakeTimers();
    const fetchMock = vi.fn(async () => {
      if (fetchMock.mock.calls.length === 1) {
        return new Response(JSON.stringify({ error: { message: "system disk overloaded" } }), {
          status: 500,
        });
      }
      return new Response(JSON.stringify({
        candidates: [{
          content: {
            parts: [{
              inlineData: {
                mimeType: "image/png",
                data: "retry-success",
              },
            }],
          },
        }],
      }), { status: 200 });
    });
    vi.stubGlobal("fetch", fetchMock);
    vi.spyOn(console, "warn").mockImplementation(() => {});

    const promise = generateWithNewApiWG("test-key", "https://newapi.example/v1", makeInput({
      model: {
        id: "apiyi_nano_banana_2",
        name: "APIYI Nano Banana 2",
        description: null,
        provider: "newapiwg",
        capabilities: ["text-to-image", "image-to-image"],
      },
      prompt: "draw a banana",
    }));

    await vi.advanceTimersByTimeAsync(500);
    const result = await promise;

    expect(fetchMock).toHaveBeenCalledTimes(2);
    expect(result.success).toBe(true);
    expect(result.outputs?.[0]).toEqual({
      type: "image",
      data: "data:image/png;base64,retry-success",
    });
  });

  it("routes sora image models through chat completions", async () => {
    const result = await generateWithNewApiWG("test-key", "https://newapi.example/v1", makeInput({
      model: {
        id: "sora_image",
        name: "sora_image",
        description: null,
        provider: "newapiwg",
        capabilities: ["text-to-image"],
      },
      prompt: "draw a banana",
    }));

    expect(result.success).toBe(true);
    expect(result.outputs?.[0]).toEqual({ type: "image", data: "data:image/png;base64,abc123" });
    expect(capturedBody!.model).toBe("sora_image");
    expect(capturedBody!.messages).toEqual([{ role: "user", content: "draw a banana" }]);
  });

  it("routes Seedream image models through images generations", async () => {
    const result = await generateWithNewApiWG("test-key", "https://newapi.example/v1", makeInput({
      model: {
        id: "seedream-5-0-260128",
        name: "seedream-5-0-260128",
        description: null,
        provider: "newapiwg",
        capabilities: ["text-to-image", "image-to-image"],
      },
      prompt: "draw a banana",
      images: ["data:image/png;base64,input"],
    }));

    expect(result.success).toBe(true);
    expect(result.outputs?.[0]).toEqual({
      type: "image",
      data: "data:image/png;base64,YWJj",
      url: "https://cdn.example.com/seedream.png",
    });
    expect(capturedBody!.model).toBe("seedream-5-0-260128");
    expect(capturedBody!.image).toBe("data:image/png;base64,input");
  });

  it("normalizes transient provider overload errors", () => {
    expect(normalizeNewApiWGError("system disk overloaded", "Nano Banana")).toBe(
      "Nano Banana: provider is temporarily overloaded. Please retry in a moment or switch models."
    );
  });

  it("returns normalized overload errors from image generation", async () => {
    vi.useFakeTimers();
    vi.stubGlobal("fetch", vi.fn(async (url: string | URL | Request, init?: RequestInit) => {
      const urlStr = typeof url === "string" ? url : url instanceof URL ? url.toString() : url.url;
      if (urlStr.includes("/images/generations") && init?.method === "POST") {
        return new Response(JSON.stringify({ error: { message: "system disk overloaded" } }), {
          status: 500,
        });
      }
      return new Response("Not Found", { status: 404 });
    }));
    vi.spyOn(console, "warn").mockImplementation(() => {});

    const promise = generateWithNewApiWG("test-key", "https://newapi.example/v1", makeInput({
      model: {
        id: "seedream-5-0-260128",
        name: "Seedream",
        description: null,
        provider: "newapiwg",
        capabilities: ["text-to-image"],
      },
      prompt: "draw a banana",
    }));

    await vi.advanceTimersByTimeAsync(2000);
    const result = await promise;

    expect(result.success).toBe(false);
    expect(result.error).toBe(
      "Seedream: provider is temporarily overloaded. Please retry in a moment or switch models."
    );
  });
});