image-studio/operations/generation/workflow.test.ts

import { describe, expect, test } from "bun:test";
import { documentCommands } from "@commands/document";
import { inpaintRegionCommands } from "@commands/inpaint-region";
import { generationCommands } from "@commands/generation";
import { commandIds } from "@commands/ids";
import { createCommandRegistry } from "@commands/registry";
import { toolCommands } from "@commands/tool";
import { createInitialAppState } from "@editor/initial-state";
import { createAppStore } from "@editor/store";
import { initialToolState } from "@editor/tools";
import type { GenerationCandidate } from "@editor/state";
import { createGenerationWorkflow, type GenerationWorkflowDependencies } from "./workflow";

describe("generation workflow", () => {
  test("reads canonical application state when generation starts", async () => {
    const app = createTestApp();
    let prompt = "";
    const workflow = createGenerationWorkflow(app.store, dependencies({
      runGenerate: async (options) => { prompt = options.settings.prompt; },
    }));

    app.store.dispatch(commandIds.toolSetGenerateSettings, { prompt: "Latest prompt" });
    await workflow.generate();

    expect(prompt).toBe("Latest prompt");
    expect(app.store.getState().editor.generation.jobs[0]?.status).toBe("succeeded");
  });

  test("uses one acceptance path to add a candidate as a layer", () => {
    const app = createTestApp();
    const artboardId = app.store.getState().document.artboards[0]?.id;
    if (!artboardId) throw new Error("Expected default artboard");
    app.store.dispatch(commandIds.generationAddCandidate, { candidate: candidate(artboardId) });
    const ids = ["asset-new", "layer-new"];
    const workflow = createGenerationWorkflow(app.store, dependencies({ createId: () => ids.shift() ?? "unused" }));

    workflow.applyCandidateAsLayer("candidate");

    expect(app.store.getState().document.assets.some((asset) => asset.id === "asset-new")).toBe(true);
    expect(app.store.getState().document.artboards[0]?.layers.some((layer) => layer.id === "layer-new")).toBe(true);
    expect(app.store.getState().editor.generation.candidates).toHaveLength(0);
  });

  test("repairs replace prerequisites by adding an aligned editable mask through commands", async () => {
    const app = createTestApp();
    const state = app.store.getState();
    state.document.assets.push({ id: "source-asset", name: "Source", mimeType: "image/png", source: "source", intrinsicSize: { w: 80, h: 60 } });
    state.document.artboards[0]!.layers.push({ id: "source", type: "raster", name: "Source", visible: true, locked: false, opacity: 1, assetId: "source-asset", transform: { position: { x: 4, y: 8 }, scale: { x: 2, y: 2 }, rotation: 0 } });
    state.editor.selection = { artboardId: "artboard", layerIds: ["source"] };
    const ids = ["mask-asset", "region"];
    const workflow = createGenerationWorkflow(app.store, dependencies({ createId: () => ids.shift() ?? "unused" }));

    await workflow.prepareInpaintMask();

    const next = app.store.getState();
    expect(next.document.inpaintRegions).toContainEqual({ id: "region", name: "Source AI edit", targetLayerId: "source", maskAssetId: "mask-asset", enabled: true });
    expect(next.editor.maskEdit).toEqual({ kind: "inpaintRegion", targetLayerId: "source", inpaintRegionId: "region", maskAssetId: "mask-asset", viewMode: "overlay" });
  });

  test("cancels the active operation and records a cancelled job", async () => {
    const app = createTestApp();
    let receivedSignal: AbortSignal | undefined;
    const workflow = createGenerationWorkflow(app.store, dependencies({
      runGenerate: async (options) => {
        receivedSignal = options.signal;
        await new Promise<void>((_resolve, reject) => options.signal?.addEventListener("abort", () => reject(new DOMException("Cancelled", "AbortError")), { once: true }));
      },
    }));

    const running = workflow.generate();
    await Promise.resolve();
    workflow.cancel();
    await running;

    expect(receivedSignal?.aborted).toBe(true);
    expect(app.store.getState().editor.generation.jobs[0]?.status).toBe("cancelled");
  });
});

function dependencies(overrides: Partial<GenerationWorkflowDependencies>): GenerationWorkflowDependencies {
  return {
    runGenerate: async () => undefined,
    runGenerateFromCandidate: async () => undefined,
    createMaskedPixelReplacementSource: async () => "replacement",
    createRefinementMask: async () => "mask",
    createInpaintRegionMask: async () => "mask",
    loadGenerationResources: async () => undefined,
    createId: () => crypto.randomUUID(),
    ...overrides,
  };
}

function candidate(artboardId: string): GenerationCandidate {
  return {
    id: "candidate",
    source: "generated",
    mimeType: "image/png",
    intrinsicSize: { w: 64, h: 64 },
    mode: "text-to-image",
    settings: initialToolState.generate,
    seed: 1,
    width: 64,
    height: 64,
    placement: {
      artboardId,
      layerName: "Generated",
      transform: { position: { x: 0, y: 0 }, scale: { x: 1, y: 1 }, rotation: 0 },
    },
  };
}

function createTestApp() {
  const state = createInitialAppState("Test");
  state.document.artboards.push({
    id: "artboard",
    name: "Artboard",
    bounds: { x: 0, y: 0, w: 100, h: 100 },
    backgroundColor: "transparent",
    visible: true,
    locked: false,
    layers: [],
  });
  const registry = createCommandRegistry([...documentCommands, ...inpaintRegionCommands, ...toolCommands, ...generationCommands]);
  return { store: createAppStore(state, registry) };
}