"""The image generation admin page, and the workflows behind it.""" from __future__ import annotations import json from pathlib import Path import pytest from fastapi.testclient import TestClient from sqlalchemy import select import lembas from lembas.db.models import ImageWorkflow from lembas.services import settings_store from lembas.services.crypto import UNCHANGED_SENTINEL, decrypt BASE_TEXT = (Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text( encoding="utf-8" ) BASE = json.loads(BASE_TEXT) @pytest.fixture def admin(client: TestClient, db, registered): from lembas.db.models import User db.query(User).update({"role": "admin"}) db.commit() return client def _settings(client, **overrides): data = { "enabled": "true", "base_url": "http://comfy.test:8188", "api_key": "", "timeout": "600", "checkpoints": "sd.safetensors\nother.safetensors", "default_workflow_id": "", "review_enabled": "true", "review_model_id": "", "max_tries": "4", "preserve_vram": "", "instructions": "", } data.update(overrides) return client.post("/admin/images", data=data, follow_redirects=False) # --- The page ------------------------------------------------------------------ def test_the_page_renders_and_is_in_the_nav(admin): page = admin.get("/admin/images").text assert 'href="/admin/images"' in page assert "ComfyUI" in page def test_settings_are_saved(admin, db): assert _settings(admin).status_code == 303 values = settings_store.images(db) assert values["enabled"] is True assert values["base_url"] == "http://comfy.test:8188" assert values["checkpoints"] == ["sd.safetensors", "other.safetensors"] assert values["review_enabled"] is True def test_a_trailing_slash_is_stripped_from_the_url(admin, db): _settings(admin, base_url="http://comfy.test:8188/") assert settings_store.images(db)["base_url"] == "http://comfy.test:8188" def test_an_absent_checkbox_is_off(admin, db): """An unticked box is simply missing from a form post -- that absence *is* the off signal.""" data = {"base_url": "http://x", "timeout": "600", "max_tries": "4"} admin.post("/admin/images", data=data, follow_redirects=False) values = settings_store.images(db) assert values["enabled"] is False assert values["review_enabled"] is False assert values["preserve_vram"] is False def test_the_numbers_are_clamped(admin, db): _settings(admin, max_tries="500", timeout="1") values = settings_store.images(db) assert values["max_tries"] == 10 assert values["timeout"] == 10.0 def test_max_tries_never_reaches_zero(admin, db): """Zero would mean the tool generates nothing and reports success. Unlike the zeroes elsewhere in the settings, there is no reading of it anybody wants.""" _settings(admin, max_tries="0") assert settings_store.images(db)["max_tries"] == 1 def test_the_key_survives_a_save_that_did_not_touch_it(admin, db): _settings(admin, api_key="secret-key") assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key" _settings(admin, api_key=UNCHANGED_SENTINEL, instructions="something else") assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key" def test_an_emptied_key_is_removed(admin, db): _settings(admin, api_key="secret-key") _settings(admin, api_key="") assert settings_store.images(db)["api_key_encrypted"] == "" def test_the_key_is_never_rendered(admin, db): _settings(admin, api_key="secret-key") assert "secret-key" not in admin.get("/admin/images").text def test_testing_without_a_url_says_so(admin, db): _settings(admin, base_url="") assert "Set a base URL first" in admin.post("/admin/images/test").text # --- Workflows ----------------------------------------------------------------- def test_new_is_not_read_as_an_id(admin): """FastAPI matches in registration order, so `/workflows/new` has to be registered before `/workflows/{id}` or "new" is an id and 404s. This has been a bug twice in this codebase.""" response = admin.get("/admin/images/workflows/new") assert response.status_code == 200 assert "Export (API)" in response.text def test_the_default_workflow_is_offered_to_start_from(admin): """Somebody setting this up for the first time should not have to find a working API-format document before they can try anything.""" page = admin.get("/admin/images/workflows/new").text assert "CheckpointLoaderSimple" in page assert "{{prompt}}" in page def test_a_workflow_is_created(admin, db): response = admin.post( "/admin/images/workflows", data={ "slug": "sdxl", "name": "SDXL", "description": "photo", "workflow": BASE_TEXT, "enabled": "true", }, follow_redirects=False, ) assert response.status_code == 303 row = db.scalar(select(ImageWorkflow).where(ImageWorkflow.slug == "sdxl")) assert row is not None assert row.workflow_json["4"]["inputs"]["ckpt_name"] == "{{model}}" def test_broken_json_is_refused_with_a_sentence(admin, db): """A sentence and not a 422: losing forty lines of pasted JSON to a validation error is not a thing to do to somebody, so the form comes back with what was typed still in it.""" response = admin.post( "/admin/images/workflows", data={"slug": "x", "name": "X", "workflow": "{not json"}, ) assert response.status_code == 200 assert "not valid JSON" in response.text assert "{not json" in response.text, "and what was typed is still there" assert db.scalar(select(ImageWorkflow)) is None def test_a_workflow_with_no_prompt_placeholder_is_refused(admin, db): """It would draw the same picture whatever anybody typed, and would look like a broken model rather than an unparameterised template.""" without = json.loads(BASE_TEXT) without["6"]["inputs"]["text"] = "a fixed prompt" response = admin.post( "/admin/images/workflows", data={"slug": "x", "name": "X", "workflow": json.dumps(without)}, ) assert "{{prompt}}" in response.text assert "every image would be the same" in response.text assert db.scalar(select(ImageWorkflow)) is None def test_an_unknown_placeholder_is_refused(admin, db): response = admin.post( "/admin/images/workflows", data={ "slug": "x", "name": "X", "workflow": json.dumps({"1": {"a": "{{prompt}} {{lora}}"}}), }, ) assert "lora" in response.text assert db.scalar(select(ImageWorkflow)) is None def test_a_ui_format_export_is_refused_helpfully(admin, db): """The two ComfyUI formats look alike enough that this is the mistake everybody makes first, and "invalid" would not tell them which button to press instead.""" response = admin.post( "/admin/images/workflows", data={"slug": "x", "name": "X", "workflow": "[]"}, ) assert "Export (API)" in response.text def test_a_duplicate_slug_is_refused(admin, db): for _ in range(2): response = admin.post( "/admin/images/workflows", data={"slug": "dup", "name": "Dup", "workflow": BASE_TEXT}, ) assert "already a workflow" in response.text assert len(list(db.scalars(select(ImageWorkflow)))) == 1 def test_a_rejected_edit_leaves_the_stored_row_alone(admin, db): admin.post( "/admin/images/workflows", data={"slug": "keep", "name": "Keep", "workflow": BASE_TEXT}, ) row = db.scalar(select(ImageWorkflow)) admin.post( f"/admin/images/workflows/{row.id}", data={"slug": "keep", "name": "Renamed", "workflow": "{broken"}, ) db.expire_all() row = db.scalar(select(ImageWorkflow)) assert row.name == "Keep", "validated against a draft, so nothing was written" def test_a_workflow_is_deleted(admin, db): admin.post( "/admin/images/workflows", data={"slug": "gone", "name": "Gone", "workflow": BASE_TEXT}, ) row = db.scalar(select(ImageWorkflow)) admin.post(f"/admin/images/workflows/{row.id}/delete", follow_redirects=False) db.expire_all() assert db.scalar(select(ImageWorkflow)) is None def test_only_an_administrator_can_reach_it(client: TestClient, db, registered): """The first account registered is an administrator, so this needs a second one -- the `test_admin_tools` fixture's reason for existing.""" from lembas.db.models import User client.post("/auth/logout") client.post( "/auth/register", data={"name": "Sam", "email": "sam@shire.test", "password": "correct horse battery"}, follow_redirects=False, ) user = db.scalar(select(User).where(User.email == "sam@shire.test")) user.role = "user" user.active = True db.commit() client.post( "/auth/login", data={"email": "sam@shire.test", "password": "correct horse battery"}, follow_redirects=False, ) assert client.get("/admin/images").status_code in (302, 303, 403, 404) assert client.post("/admin/images/workflows", data={}).status_code in (302, 303, 403, 404)