You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

337 lines
12 KiB

import { describe, expect, it, beforeEach, vi } from "vitest";
import { useGenerationPreferenceStore, createInitialGenerationPreferences } from "@/store/generationPreferenceStore";
import { useModelStore } from "@/store/modelStore";
import { buildDefaultGenerationNodeCreation, buildDefaultGenerationNodeData } from "@/utils/defaultGenerationNodeModel";
import {
DEFAULT_AUDIO_GENERATION_MODEL_NAME,
DEFAULT_IMAGE_GENERATION_ASPECT_RATIO,
DEFAULT_IMAGE_GENERATION_MODEL_NAME,
DEFAULT_LLM_GENERATION_MODEL_NAME,
DEFAULT_VIDEO_GENERATION_ASPECT_RATIO,
DEFAULT_VIDEO_GENERATION_MODEL_NAME,
} from "@/constants/generationPreferenceDefaults";
import type { SelectedModel } from "@/types";
import type { ProviderModel } from "@/lib/providers/types";
const selectedModel: SelectedModel = {
provider: "popiserver",
modelId: "image-model",
displayName: "Image Model",
subType: 102,
subTypes: [102],
capabilities: ["text-to-image"],
};
function providerModel(id: string, name: string, subType: number, capabilities: ProviderModel["capabilities"]): ProviderModel {
return {
id,
name,
description: null,
provider: "popiserver",
capabilities,
metadata: {
subType,
subTypes: [subType],
},
};
}
type GenerationPreferenceState = ReturnType<typeof useGenerationPreferenceStore.getState>;
type ImagePreferencePatch = Parameters<GenerationPreferenceState["updateImagePreference"]>[0];
describe("generationPreferenceStore", () => {
beforeEach(() => {
vi.restoreAllMocks();
useGenerationPreferenceStore.setState(createInitialGenerationPreferences());
useModelStore.getState().setPopiModels([
providerModel("image-model", "Image Model", 102, ["text-to-image"]),
providerModel("video-model", "Video Model", 202, ["text-to-video"]),
providerModel("existing-model-id", "Existing Model", 202, ["text-to-video"]),
]);
vi.spyOn(globalThis, "fetch").mockResolvedValue({
ok: true,
json: async () => ({ success: true, models: [] }),
} as Response);
});
it("initializes image and video preferences from default generation config", async () => {
const { image, video } = useGenerationPreferenceStore.getState();
expect(image.model.modelId).toBe("");
expect(image.model.displayName).toBe(DEFAULT_IMAGE_GENERATION_MODEL_NAME);
expect(image.aspectRatio).toBe(DEFAULT_IMAGE_GENERATION_ASPECT_RATIO);
expect(image.parameters).toEqual({ ratio: DEFAULT_IMAGE_GENERATION_ASPECT_RATIO });
expect(video.model.modelId).toBe("");
expect(video.model.displayName).toBe(DEFAULT_VIDEO_GENERATION_MODEL_NAME);
expect(video.parameters).toEqual({ ratio: DEFAULT_VIDEO_GENERATION_ASPECT_RATIO });
expect(useGenerationPreferenceStore.getState().audio.model.modelId).toBe("");
expect(useGenerationPreferenceStore.getState().audio.model.displayName).toBe(DEFAULT_AUDIO_GENERATION_MODEL_NAME);
});
it("keeps only model configuration fields when remembering image preferences", async () => {
useGenerationPreferenceStore.getState().updateImagePreference({
selectedModel,
batchSize: 3,
aspectRatio: "16:9",
resolution: "2K",
parameters: {
ratio: "16:9",
images: ["parameter-image-should-stay-if-schema-owns-it"],
type: 1,
},
extraTaskParams: { generate_audio: false },
subType: 101,
prompt: "do not remember prompts",
inputImages: ["https://example.com/reference.png"],
inputVideos: ["https://example.com/reference.mp4"],
voices: ["https://example.com/reference.wav"],
referenceSubjectList: [{ id: "1", type: "image", name: "Reference" }],
} as ImagePreferencePatch);
const imagePreference = useGenerationPreferenceStore.getState().image as unknown as Record<string, unknown>;
expect(imagePreference.model).toEqual(selectedModel);
expect(imagePreference.batchSize).toBe(3);
expect(imagePreference.aspectRatio).toBe("16:9");
expect(imagePreference.resolution).toBe("2K");
expect(imagePreference.parameters).toEqual({
ratio: "16:9",
images: ["parameter-image-should-stay-if-schema-owns-it"],
type: 1,
});
expect(imagePreference.extraTaskParams).toEqual({ generate_audio: false });
expect(imagePreference.subType).toBe(101);
expect(imagePreference).not.toHaveProperty("prompt");
expect(imagePreference).not.toHaveProperty("inputImages");
expect(imagePreference).not.toHaveProperty("inputVideos");
expect(imagePreference).not.toHaveProperty("voices");
expect(imagePreference).not.toHaveProperty("referenceSubjectList");
expect(imagePreference).not.toHaveProperty("nodeSize");
const newNodeData = await buildDefaultGenerationNodeData("nanoBanana");
expect(newNodeData).toMatchObject({
selectedModel,
batchSize: 3,
aspectRatio: "16:9",
resolution: "2K",
parameters: {
ratio: "16:9",
images: ["parameter-image-should-stay-if-schema-owns-it"],
},
extraTaskParams: { generate_audio: false },
subType: 101,
});
expect(newNodeData).not.toHaveProperty("prompt");
expect(newNodeData).not.toHaveProperty("inputImages");
expect(newNodeData).not.toHaveProperty("inputVideos");
expect(newNodeData).not.toHaveProperty("voices");
expect(newNodeData).not.toHaveProperty("referenceSubjectList");
await expect(buildDefaultGenerationNodeCreation("nanoBanana")).resolves.toEqual({
data: expect.any(Object),
size: { width: 615, height: 346 },
});
});
it("uses configured video ratio when creating generation nodes from default preferences", async () => {
useGenerationPreferenceStore.getState().updateVideoPreference({
selectedModel: {
provider: "popiserver",
modelId: "video-model",
displayName: "Video Model",
subType: 202,
subTypes: [202],
capabilities: ["text-to-video"],
},
parameters: { ratio: DEFAULT_VIDEO_GENERATION_ASPECT_RATIO },
});
await expect(buildDefaultGenerationNodeCreation("generateVideo")).resolves.toEqual({
data: expect.any(Object),
size: { width: 615, height: 346 },
});
});
it("uses parameter ratio to update image aspect ratio preference", async () => {
useGenerationPreferenceStore.getState().updateImagePreference({
selectedModel,
aspectRatio: "16:9",
parameters: { ratio: "4:5" },
} as ImagePreferencePatch);
const imagePreference = useGenerationPreferenceStore.getState().image;
expect(imagePreference.aspectRatio).toBe("4:5");
await expect(buildDefaultGenerationNodeCreation("nanoBanana")).resolves.toEqual({
data: expect.objectContaining({
aspectRatio: "4:5",
parameters: { ratio: "4:5" },
}),
size: { width: 346, height: 432 },
});
});
it("does not remember a model when its subtype is outside the target model store kind", () => {
useModelStore.getState().setPopiModels([
providerModel("regular-image", "Regular Image", 102, ["text-to-image"]),
providerModel("multi-angle", "Multi Angle", 113, ["image-to-image"]),
]);
const previousImagePreference = useGenerationPreferenceStore.getState().image;
useGenerationPreferenceStore.getState().updateImagePreference({
selectedModel: {
provider: "popiserver",
modelId: "multi-angle",
displayName: "Multi Angle",
subType: 113,
subTypes: [113],
capabilities: ["image-to-image"],
},
parameters: { ratio: "9:16" },
aspectRatio: "9:16",
});
expect(useGenerationPreferenceStore.getState().image).toEqual(previousImagePreference);
});
it("resolves only missing model id from stored model name and keeps the stored parameters", async () => {
const videoModel: ProviderModel = {
id: "44",
name: "Seedance",
description: null,
provider: "popiserver",
capabilities: ["text-to-video", "image-to-video"],
metadata: {
billingMethod: 2,
aiModelId: 9001,
aiModelCode: "seedance-video-pro",
aiModelCodeAlias: "seedance-pro",
aiModelName: "seedance 2.0-fast",
type: 2,
subType: 101,
subTypes: [101, 102],
},
};
useModelStore.getState().setPopiModels([videoModel]);
vi.mocked(globalThis.fetch).mockResolvedValueOnce({
ok: true,
json: async () => ({ success: true, models: [videoModel] }),
} as Response);
useGenerationPreferenceStore.getState().updateVideoPreference({
aiModelName: "seedance 2.0-fast",
parameters: { ratio: "9:16" },
batchSize: 2,
});
const newNodeData = await buildDefaultGenerationNodeData("generateVideo");
expect(newNodeData).toMatchObject({
batchSize: 2,
parameters: { ratio: "9:16" },
selectedModel: {
provider: "popiserver",
modelId: "44",
displayName: "seedance 2.0-fast",
billingMethod: 2,
},
});
expect(useGenerationPreferenceStore.getState().video.model.modelId).toBe("44");
expect(useGenerationPreferenceStore.getState().video.model.displayName).toBe("seedance 2.0-fast");
expect(useGenerationPreferenceStore.getState().video.model.billingMethod).toBe(2);
await expect(buildDefaultGenerationNodeCreation("generateVideo")).resolves.toEqual({
data: expect.objectContaining({
parameters: { ratio: "9:16" },
selectedModel: expect.objectContaining({ modelId: "44", billingMethod: 2 }),
}),
size: { width: 346, height: 615 },
});
});
it("does not look up a default model when the preference already has a model id", async () => {
useGenerationPreferenceStore.getState().updateVideoPreference({
selectedModel: {
provider: "popiserver",
modelId: "existing-model-id",
displayName: "Existing Model",
subType: 202,
subTypes: [202],
},
parameters: { ratio: "16:9" },
});
const newNodeData = await buildDefaultGenerationNodeData("generateVideo");
expect(newNodeData).toMatchObject({
parameters: { ratio: "16:9" },
selectedModel: {
provider: "popiserver",
modelId: "existing-model-id",
displayName: "Existing Model",
},
});
expect(globalThis.fetch).not.toHaveBeenCalled();
});
it("uses the same name-to-id lookup for default image, video, audio, and llm models", async () => {
const models: ProviderModel[] = [
{
id: "img-32",
name: DEFAULT_IMAGE_GENERATION_MODEL_NAME,
description: null,
provider: "popiserver",
capabilities: ["text-to-image"],
metadata: { billingMethod: 2, subType: 102, subTypes: [102] },
},
{
id: "vid-44",
name: DEFAULT_VIDEO_GENERATION_MODEL_NAME,
description: null,
provider: "popiserver",
capabilities: ["text-to-video"],
metadata: { billingMethod: 1, subType: 202, subTypes: [202] },
},
{
id: "aud-55",
name: DEFAULT_AUDIO_GENERATION_MODEL_NAME,
description: null,
provider: "popiserver",
capabilities: ["text-to-audio"],
metadata: { subType: 301, subTypes: [301] },
},
{
id: "llm-66",
name: DEFAULT_LLM_GENERATION_MODEL_NAME,
description: null,
provider: "popiserver",
capabilities: ["text-to-text"],
metadata: { type: 5 },
},
];
vi.mocked(globalThis.fetch).mockResolvedValueOnce({
ok: true,
json: async () => ({ success: true, models }),
} as Response);
useModelStore.getState().setPopiModels(models);
const imageData = await buildDefaultGenerationNodeData("nanoBanana");
const videoData = await buildDefaultGenerationNodeData("generateVideo");
const audioData = await buildDefaultGenerationNodeData("generateAudio");
const llmData = await buildDefaultGenerationNodeData("llmGenerate");
expect(imageData).toMatchObject({
selectedModel: { modelId: "img-32", displayName: DEFAULT_IMAGE_GENERATION_MODEL_NAME, billingMethod: 2 },
});
expect(videoData).toMatchObject({
selectedModel: { modelId: "vid-44", displayName: DEFAULT_VIDEO_GENERATION_MODEL_NAME, billingMethod: 1 },
});
expect(audioData).toMatchObject({
selectedModel: { modelId: "aud-55", displayName: DEFAULT_AUDIO_GENERATION_MODEL_NAME },
});
expect(llmData).toMatchObject({
model: "llm-66",
selectedModel: { modelId: "llm-66", displayName: DEFAULT_LLM_GENERATION_MODEL_NAME },
});
});
});