ComfyAssets-ComfyUI_PromptManager/tests/browser/workspace_collections_fixture.py

"""Owned, deterministic collections for the workspace collection workflows (#331).

Twenty-one PNGs under the output source: twelve listed members of
"Landscapes", eight of "Experiments" and one loose file in neither. Landscapes
also holds one member whose file does not exist (MC3: hidden, counted as
missing). One saved view reopens "Experiments" (D140). Members are added
through ``ImageCollectionOperations`` with favorites' keys
(``_favorite_root_key``, ``_favorite_relative_key``), as the routes add them.
"""

from __future__ import annotations

import os
import time
from pathlib import Path
from typing import TYPE_CHECKING

from aiohttp import web
from PIL import Image
from tests.browser.folder_trash import replace_folder

if TYPE_CHECKING:
    from py.api import PromptManagerAPI

FOLDER = "collections"
LANDSCAPES = "Landscapes"
EXPERIMENTS = "Experiments"
LANDSCAPE_FILES = tuple(f"landscape_{index:02d}.png" for index in range(1, 13))
EXPERIMENT_FILES = tuple(f"experiment_{index:02d}.png" for index in range(1, 9))
LOOSE_FILE = "loose.png"
MISSING_FILE = "gone.png"
SAVED_VIEW = "Experiments view"
_TABLES = (
    "image_collection_items",
    "image_collections",
    "saved_views",
    "image_favorites",
    "media_review_decisions",
    "media_review_items",
    "media_image_tags",
    "media_annotation_bytes",
    "media_annotation_files",
    "image_prompt_pairs",
    "prompt_pair_tags",
    "prompt_pairs",
    "generated_images",
    "prompt_tags",
    "tags",
    "prompt_library_state",
    "prompts",
)
_IMAGES_STATE = {
    "q": "",
    "category": "",
    "sort": "newest",
    "source": "all",
    "folder": "",
    "media": "all",
    "tags": [],
    "favorites": False,
}


def _write_png(path: Path, shade: int, stamp: float) -> None:
    with Image.new("RGB", (64, 48), (shade, 120, 255 - shade)) as image:
        image.save(path)
    os.utime(path, (stamp, stamp))


def seed_collections(api: PromptManagerAPI, root: Path) -> dict:
    """Replace this server's disposable rows and files with the two collections."""
    # The API package needs the host stubs the server installs before importing it.
    from py.api.images import _favorite_relative_key, _favorite_root_key

    db = api.db
    with db.model.get_connection() as conn:
        for table in _TABLES:
            conn.execute(f"DELETE FROM {table}")
    for name in ("output", "input"):
        replace_folder(root / name)
    folder = root / "output" / FOLDER
    folder.mkdir()
    # Newest first in this order: Landscapes, then Experiments, then the loose file.
    now = time.time()
    ordered = (*LANDSCAPE_FILES, *EXPERIMENT_FILES, LOOSE_FILE)
    for index, name in enumerate(ordered):
        _write_png(folder / name, (index * 11) % 256, now - 60 - index * 60)
    root_key = _favorite_root_key(list(api._get_all_output_dirs())[0])
    members = {
        LANDSCAPES: (*LANDSCAPE_FILES, MISSING_FILE),
        EXPERIMENTS: EXPERIMENT_FILES,
    }
    ids = {}
    for name, files in members.items():
        added = db.add_to_collection(
            name=name,
            items=[
                (root_key, _favorite_relative_key(f"{FOLDER}/{file}")) for file in files
            ],
        )
        ids[name] = added["collection"]["id"]
    view = db.create_saved_view(
        "images", SAVED_VIEW, {**_IMAGES_STATE, "collection": str(ids[EXPERIMENTS])}
    )
    api.invalidate_gallery_cache()
    return {
        "collections": {
            LANDSCAPES: {"id": ids[LANDSCAPES], "count": 12, "missing": 1},
            EXPERIMENTS: {"id": ids[EXPERIMENTS], "count": 8, "missing": 0},
        },
        "files": {
            LANDSCAPES: [f"{FOLDER}/{file}" for file in LANDSCAPE_FILES],
            EXPERIMENTS: [f"{FOLDER}/{file}" for file in EXPERIMENT_FILES],
            "loose": f"{FOLDER}/{LOOSE_FILE}",
        },
        "saved_view": {"id": view["id"], "name": SAVED_VIEW},
    }


def register_collections_commands(
    routes: web.RouteTableDef, api: PromptManagerAPI, root: Path
) -> None:
    """Register the fixed seed command; no path or SQL comes from a request."""

    @routes.post("/__test__/collections/seed")
    async def seed(request: web.Request) -> web.Response:
        data = await request.json() if request.can_read_body else {}
        if data != {} or request.query:
            raise web.HTTPBadRequest()
        manifest = await api._run_in_executor(seed_collections, api, root)
        return web.json_response({"success": True, **manifest})