olivv-cs-ComfyUI-FunPack/core/tests/test_addons.py

"""Add-ons: what is offered, what a switch does, and what the next start installs."""

import asyncio
import json

import pytest

from core import addons, config, update


@pytest.fixture
def here(tmp_path, monkeypatch):
    """addons.json in a temp folder; `present` is the set of add-on ids whose packages 'are installed'."""
    monkeypatch.setattr(config, "ADDONS_FILE", tmp_path / "addons.json")
    present = set()
    monkeypatch.setattr(addons, "_present", lambda a: a["id"] in present)
    for a in addons.ADDONS:
        monkeypatch.setitem(a, "offered", lambda: True)
    return present


def state(aid):
    return next(a["state"] for a in addons.listing() if a["id"] == aid)


def test_nothing_is_wanted_until_asked(here):
    assert state("nvidia_vsr") == "available" and state("camera") == "available"
    assert not addons.enabled("camera")


def test_what_was_installed_before_add_ons_existed_stays_on(here):
    here.add("camera")
    assert state("camera") == "installed" and addons.enabled("camera")


def test_switching_on_waits_for_the_restart_and_off_keeps_the_packages(here):
    assert addons.set_wanted("nvidia_vsr", True)
    assert state("nvidia_vsr") == "pending" and not addons.enabled("nvidia_vsr")
    here.add("nvidia_vsr")                                   # the restart installed it
    assert state("nvidia_vsr") == "installed" and addons.enabled("nvidia_vsr")
    addons.set_wanted("nvidia_vsr", False)
    assert state("nvidia_vsr") == "disabled" and not addons.enabled("nvidia_vsr")


def test_an_add_on_this_machine_cannot_run_is_not_offered(here, monkeypatch):
    for gone in ("nvidia_vsr", "nvidia_framegen"):
        monkeypatch.setitem(addons._BY_ID[gone], "offered", lambda: False)
    assert [a["id"] for a in addons.listing()] == ["camera", "lora_training"]
    assert addons.set_wanted("nvidia_vsr", True) is False
    assert addons.set_wanted("no-such-addon", True) is False


def test_the_start_installs_only_what_is_wanted_and_missing(here, monkeypatch):
    ran = []

    def fake(lines, upgrade=()):
        ran.append(list(lines))
        here.add("nvidia_vsr")
        return {"ok": True, "detail": ""}
    monkeypatch.setattr(update, "install_lines", fake)
    addons.install_wanted()
    assert ran == [], "installed something nobody asked for"
    addons.set_wanted("nvidia_vsr", True)
    addons.install_wanted()
    assert ran == [["nvidia-vfx>=0.2"]]
    addons.install_wanted()
    assert len(ran) == 1, "installed again though it is already here"


def test_a_failed_install_is_said_and_not_retried_until_switched_off_and_on(here, monkeypatch):
    ran = []
    monkeypatch.setattr(update, "install_lines", lambda lines, upgrade=(): ran.append(1) or {"ok": False, "detail": "no wheel for this Python"})
    addons.set_wanted("nvidia_vsr", True)
    addons.install_wanted()
    row = next(a for a in addons.listing() if a["id"] == "nvidia_vsr")
    assert row["failed"] == "no wheel for this Python" and row["state"] == "available", "a failed install must not read as waiting for a restart"
    addons.install_wanted()
    assert len(ran) == 1, "a broken install would be retried on every start"
    addons.set_wanted("nvidia_vsr", False)
    addons.set_wanted("nvidia_vsr", True)
    addons.install_wanted()
    assert len(ran) == 2


def test_install_lines_skips_what_is_present_and_moves_nothing(monkeypatch):
    seen = {}
    monkeypatch.setattr(update, "_installed_version", lambda n: "1" if n == "gguf" else None)
    monkeypatch.setattr(update, "_pip_freeze", lambda: {"torch": "2.10.0"})
    monkeypatch.setattr(update, "_install_each", lambda status, path, before, note, timeout, broken:
                        seen.update(lines=status["install"], pins=open(path).read()) or {"ok": True})
    update.install_lines(["gguf", "nvidia-vfx>=0.2"])
    assert seen["lines"] == ["nvidia-vfx>=0.2"]
    assert "torch==2.10.0" in seen["pins"]


def test_the_routes_list_and_switch(comfyui, here):
    from aiohttp import web
    from aiohttp.test_utils import TestClient, TestServer
    from core import routes

    async def go():
        app, table = web.Application(), web.RouteTableDef()
        routes.register(table, prefix="/fun")
        app.add_routes(table)
        async with TestClient(TestServer(app)) as client:
            got = await (await client.get("/fun/api/addons")).json()
            bad = await client.post("/fun/api/addons", json={"id": "camera", "enabled": "yes"})
            ok = await (await client.post("/fun/api/addons", json={"id": "camera", "enabled": True})).json()
            return got, bad.status, ok
    got, bad, ok = asyncio.run(go())
    assert {a["id"] for a in got["addons"]} == {"nvidia_vsr", "nvidia_framegen", "camera", "lora_training"}
    assert bad == 400
    assert next(a["state"] for a in ok["addons"] if a["id"] == "camera") == "pending"
    assert json.loads(config.ADDONS_FILE.read_text())["wanted"] == ["camera"]


def test_an_add_on_this_machine_cannot_run_is_not_enabled_even_if_wanted_and_installed(here, monkeypatch):
    here.add("nvidia_vsr")
    addons.set_wanted("nvidia_vsr", True)
    assert addons.enabled("nvidia_vsr")
    monkeypatch.setitem(addons._BY_ID["nvidia_vsr"], "offered", lambda: False)     # the folder moved to a machine without the card
    assert not addons.enabled("nvidia_vsr")


def test_lora_training_starts_on_when_its_package_is_already_there(here):
    here.add("lora_training")
    assert addons.enabled("lora_training")
    assert next(a["state"] for a in addons.listing() if a["id"] == "lora_training") == "installed"
    addons.set_wanted("lora_training", False)
    assert not addons.enabled("lora_training")


def test_the_camera_add_on_pins_the_preshed_that_can_load_spacy_models():
    camera = next(a for a in addons.ADDONS if a["id"] == "camera")
    assert addons.PRESHED_PIN in camera["requires"]


def test_preshed_3_0_14_is_flagged_and_older_is_not(monkeypatch):
    from importlib import metadata
    for have, bad in (("3.0.14", True), ("3.0.13", False), ("3.1.0", True)):
        monkeypatch.setattr(metadata, "version", lambda _n, have=have: have)
        assert addons.preshed_broken() is bad


def test_frame_generation_is_a_separate_switch_on_the_same_package():
    vsr, fg = addons._BY_ID["nvidia_vsr"], addons._BY_ID["nvidia_framegen"]
    assert vsr["requires"] == fg["requires"] == ("nvidia-vfx>=0.2",)       # neither add-on can ask for the 0.1 that has no frame generation
    assert fg["floors"] == {"nvidia-vfx": "0.2"} and "floors" not in vsr


def test_an_old_nvidia_vfx_is_present_for_super_resolution_but_not_for_frame_generation(monkeypatch):
    monkeypatch.setattr(addons.importlib.util, "find_spec", lambda m: object())
    for have, ok in (("0.1.0.1", False), ("0.2.0.0", True), ("0.10.0", True)):
        monkeypatch.setattr(update, "_installed_version", lambda n, have=have: have)
        assert addons._present(addons._BY_ID["nvidia_framegen"]) is ok
        assert addons._present(addons._BY_ID["nvidia_vsr"]) is True


def _pip(monkeypatch, installed, freeze):
    seen = {}
    monkeypatch.setattr(update, "_installed_version", lambda n: installed.get(n))
    monkeypatch.setattr(update, "_pip_freeze", lambda: freeze)
    monkeypatch.setattr(update, "_install_each", lambda status, path, before, note, timeout, broken:
                        seen.update(lines=status["install"], pins=open(path).read()) or {"ok": True})
    return seen


def test_frame_generation_moves_an_old_nvidia_vfx_up_and_nothing_else(monkeypatch):
    seen = _pip(monkeypatch, {"nvidia-vfx": "0.1.0.1"}, {"torch": "2.10.0", "nvidia-vfx": "0.1.0.1"})
    update.install_lines(["nvidia-vfx>=0.2"], upgrade=("nvidia-vfx",))
    assert seen["lines"] == ["nvidia-vfx>=0.2"]
    assert "torch==2.10.0" in seen["pins"] and "nvidia-vfx" not in seen["pins"]


def test_super_resolution_never_moves_a_newer_nvidia_vfx_back(monkeypatch):
    seen = _pip(monkeypatch, {"nvidia-vfx": "0.2.0.0"}, {"nvidia-vfx": "0.2.0.0"})
    assert update.install_lines(["nvidia-vfx>=0.2"])["ran"] is False and not seen          # present: nothing to do, in either order
    assert update.install_lines(["nvidia-vfx>=0.2"], upgrade=("nvidia-vfx",))["ran"] is False   # already at the floor