ComfyAssets-ComfyUI_PromptManager/tests/browser/workspace_integrations_fixture.py

"""System › Integrations test hooks for the browser tests (#249 System S10).

A fake LoraManager install lives under this server's temp root; nothing ever
reaches a real LoraManager, ComfyUI install or the network:

- ``POST /__test__/integrations/reset`` turns the integration off (no path,
  no key, no trigger words), removes the fake install, its LoRA folder and
  the image cache, forgets every LoRA import and injected failure, and opens
  the download gate;
- ``POST /__test__/integrations/lora-manager`` ``{loras, civitai, previous}``
  writes ``custom_nodes/ComfyUI-Lora-Manager`` (detected through the
  ``folder_paths`` stub's base path) and ``loras`` metadata files under
  ``models/loras``, each with a local preview PNG; ``civitai`` adds one
  CivitAI example image per LoRA (served by the stub opener below);
  ``previous`` saves that many earlier ``lora-manager`` imports with tags and
  image links;
- ``POST /__test__/integrations/gate`` ``{hold}`` holds every CivitAI
  download until it is released (``{"hold": false}``), so a test sees the job
  while it downloads;
- ``POST /__test__/integrations/fail-import`` makes the next collection
  replacement fail (a database error after staging);
- ``POST /__test__/integrations/pair`` saves a pair whose positive is the first
  imported prompt (a saved pair that blocks the replacement);
- ``GET /__test__/integrations/state`` answers the imported prompt texts,
  whether a CivitAI key is saved and the downloads served.

CivitAI downloads are always answered by a stub opener (a small PNG), so no
test ever opens a network connection.
"""

import io
import json
import shutil
import threading
from pathlib import Path

from aiohttp import web
from PIL import Image

MAX_LORAS = 20
CATEGORY = "lora-manager"


def _png():
    buffer = io.BytesIO()
    Image.new("RGB", (8, 8), "#3a6ea5").save(buffer, "PNG")
    return buffer.getvalue()


class _Response:
    def __init__(self, body):
        self._body = io.BytesIO(body)
        self.headers = {"Content-Length": str(len(body))}

    def read(self, amount=-1):
        return self._body.read(amount)

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        return False


def register_integrations_commands(app, routes, api, root: Path):
    from py import lora_utils
    from py.config import IntegrationConfig
    from utils.hashing import generate_prompt_hash

    lora_manager = root / "custom_nodes" / "ComfyUI-Lora-Manager"
    loras = root / "models" / "loras"
    gate = threading.Event()
    gate.set()
    served = []
    failures = {"replace": 0}
    pairs = []  # the blocking pairs this fixture saved
    real_replace = type(api.db).replace_category_prompts

    def stub_open(request, timeout=10):
        served.append(request.full_url)
        gate.wait(30)
        return _Response(_png())

    lora_utils._open_url = stub_open

    def replace(db, *args, **kwargs):
        if failures["replace"] > 0:
            failures["replace"] -= 1
            raise RuntimeError("Injected replacement failure")
        return real_replace(db, *args, **kwargs)

    type(api.db).replace_category_prompts = replace

    def reset():
        IntegrationConfig.LORA_MANAGER_ENABLED = False
        IntegrationConfig.LORA_MANAGER_PATH = ""
        IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = False
        IntegrationConfig.CIVITAI_API_KEY = ""
        lora_utils.get_trigger_cache().clear()
        gate.set()
        served.clear()
        failures["replace"] = 0
        cache = Path(lora_utils.get_lora_image_cache_dir(create=False))
        for folder in (lora_manager, loras, cache):
            shutil.rmtree(folder, ignore_errors=True)
        for pair_id in pairs:
            api.db.delete_prompt_pair(pair_id)
        pairs.clear()
        real_replace(api.db, CATEGORY, [], allow_empty=True)

    @routes.post("/__test__/integrations/reset")
    async def integrations_reset(request):
        reset()
        return web.json_response({"success": True})

    @routes.post("/__test__/integrations/lora-manager")
    async def make_lora_manager(request):
        data = await request.json() if request.can_read_body else {}
        count = data.get("loras", 2)
        previous = data.get("previous", 0)
        for value in (count, previous):
            if isinstance(value, bool) or not isinstance(value, int):
                raise web.HTTPBadRequest()
        if not 0 <= count <= MAX_LORAS or not 0 <= previous <= MAX_LORAS:
            raise web.HTTPBadRequest()
        (lora_manager / "py").mkdir(parents=True, exist_ok=True)
        (lora_manager / "__init__.py").write_text("")
        loras.mkdir(parents=True, exist_ok=True)
        for number in range(count):
            stem = f"style_{number + 1:02}"
            civitai = {"trainedWords": [f"trigger{number + 1}"]}
            if data.get("civitai") is True:
                civitai["images"] = [
                    {
                        "url": f"https://image.civitai.com/x/{stem}.jpeg",
                        "meta": {"prompt": f"{stem} example prompt"},
                    }
                ]
            meta = {
                "file_name": f"{stem}.safetensors",
                "model_name": f"Style {number + 1:02}",
                "civitai": civitai,
            }
            (loras / f"{stem}.metadata.json").write_text(json.dumps(meta))
            (loras / f"{stem}.png").write_bytes(_png())
        for number in range(previous):
            text = f"previous lora import {number + 1}"
            prompt_id = api.db.save_prompt(
                text=text,
                category=CATEGORY,
                tags=[CATEGORY, f"lora:previous-{number + 1}"],
                prompt_hash=generate_prompt_hash(text),
            )
            api.db.link_image_to_prompt(str(prompt_id), str(loras / "old.png"))
        return web.json_response({"success": True})

    @routes.post("/__test__/integrations/gate")
    async def download_gate(request):
        data = await request.json()
        if not isinstance(data.get("hold"), bool):
            raise web.HTTPBadRequest()
        if data["hold"]:
            gate.clear()
        else:
            gate.set()
        return web.json_response({"success": True})

    @routes.post("/__test__/integrations/fail-import")
    async def fail_import(request):
        failures["replace"] = 1
        return web.json_response({"success": True})

    @routes.post("/__test__/integrations/pair")
    async def pair(request):
        with api.db.model.get_connection() as conn:
            row = conn.execute(
                "SELECT text FROM prompts WHERE category = ? ORDER BY id LIMIT 1",
                (CATEGORY,),
            ).fetchone()
        if row is None:
            raise web.HTTPConflict()
        pairs.append(api.db.save_prompt_pair(row[0], "blurry"))
        return web.json_response({"success": True})

    @routes.get("/__test__/integrations/state")
    async def state(request):
        imports = api.db.search_prompts(category=CATEGORY, limit=MAX_LORAS * 2)
        return web.json_response(
            {
                "imports": sorted(prompt["text"] for prompt in imports),
                "key_saved": bool(IntegrationConfig.CIVITAI_API_KEY),
                "enabled": IntegrationConfig.LORA_MANAGER_ENABLED,
                "trigger_words": IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED,
                "downloads": len(served),
            }
        )

    return reset