import { expect, test } from "@playwright/test"; import { createAtlasCloudImagePreset, inferAtlasCloudMediaKind, } from "@/lib/ai/atlascloud-models"; import { createAtlasCloudRequestBody } from "@/lib/ai/atlascloud-request"; import { is3DModelDescriptor, isImageGenerationModelDescriptor, isVideoModelDescriptor, } from "@/lib/ai/model-classification"; const atlasHunyuan3D = { model: "tencent/hunyuan3d-pro/text-to-3d", type: "Image", displayName: "Hunyuan 3D Pro Text-to-3D", categories: ["TEXT-TO-3D"], tags: ["TEXT-TO-3D", "NEW"], }; test.describe("AI model media classification", () => { test("keeps multi-reference BananaImage models classified as image models", () => { const descriptor = { providerType: "bananaimage", modelId: "gemini-3-pro-image-preview", capabilities: { supportedModes: [ "text-to-image", "image-to-image", "reference-images", ], }, supportsImageGen: true, }; expect(isVideoModelDescriptor(descriptor)).toBe(false); expect(isImageGenerationModelDescriptor(descriptor)).toBe(true); }); test("still recognizes explicit video generation modes", () => { expect( isVideoModelDescriptor({ providerType: "custom", modelId: "custom-video-model", capabilities: { supportedModes: ["text-to-video"] }, }) ).toBe(true); }); test("keeps Atlas 3D entries out of image and video categories", () => { expect(inferAtlasCloudMediaKind(atlasHunyuan3D)).toBe("3d"); const descriptor = { providerType: "atlascloud", modelId: atlasHunyuan3D.model, capabilities: { kind: "3d", mediaType: "3d", supportedModes: ["text-to-3d"], }, supportsImageGen: false, }; expect(is3DModelDescriptor(descriptor)).toBe(true); expect(isImageGenerationModelDescriptor(descriptor)).toBe(false); expect(isVideoModelDescriptor(descriptor)).toBe(false); }); test("writes 3D Atlas presets with model-level capabilities", () => { const preset = createAtlasCloudImagePreset( atlasHunyuan3D, { components: { schemas: { Input: { type: "object", properties: { model: { type: "string" }, prompt: { type: "string" }, }, required: ["model", "prompt"], }, }, }, }, 0 ); expect(preset.supportsImageGen).toBe(false); expect(preset.capabilities).toMatchObject({ kind: "3d", mediaType: "3d", supportedModes: ["text-to-3d"], }); }); test("still distinguishes Atlas image and video descriptors", () => { const imageDescriptor = { providerType: "atlascloud", modelId: "black-forest-labs/flux-2-pro/text-to-image", capabilities: { kind: "text-to-image", mediaType: "image", supportedModes: ["text-to-image"], }, supportsImageGen: true, }; const videoDescriptor = { providerType: "atlascloud", modelId: "veed/fabric-1.0/image-to-video", capabilities: { kind: "video", mediaType: "video", supportedModes: ["image-to-video"], }, supportsImageGen: false, }; expect(isImageGenerationModelDescriptor(imageDescriptor)).toBe(true); expect(isVideoModelDescriptor(imageDescriptor)).toBe(false); expect(isVideoModelDescriptor(videoDescriptor)).toBe(true); expect(isImageGenerationModelDescriptor(videoDescriptor)).toBe(false); }); test("maps Atlas reference-video image models to a valid frontend request", () => { const preset = createAtlasCloudImagePreset( { model: "google/nano-banana-2/reference-to-image", type: "Image", displayName: "Nano Banana 2 Reference to Image", categories: ["IMAGE-TO-IMAGE"], }, { paths: { "/api/v1/model/generateImage": { post: {} }, "/api/v1/model/prediction/{request_id}": { get: {} }, }, components: { schemas: { Input: { type: "object", properties: { model: { type: "string" }, prompt: { type: "string" }, resolution: { type: "string", enum: ["1k", "2k", "4k"] }, video_clips: { type: "array", minItems: 1, maxItems: 1, items: { type: "object", properties: { url: { type: "string" }, start: { type: "number", default: 0 }, ends: { type: "number", default: 0 }, fps: { type: "number", default: 1 }, }, }, }, }, required: ["model", "prompt", "video_clips"], }, }, }, } ); expect(preset.capabilities).toMatchObject({ mediaType: "image", supportsVideoReference: true, requiresVideoReference: true, maxVideoReferences: 1, }); const { body } = createAtlasCloudRequestBody({ modelId: preset.modelId, prompt: "参考视频生成商品图", referenceImages: [], referenceVideos: ["https://example.com/reference.mp4"], imageParams: { imageSize: "2K", aspectRatio: "1:1" }, modelConfig: preset.imageGenConfig, }); expect(body.video_clips).toEqual([ { url: "https://example.com/reference.mp4", start: 0, ends: 0, fps: 1, }, ]); }); });