"""System test hooks for the browser tests (#249 System S2 Task 6). The data is the concept-shaped real fixture (``POST /__test__/concept/reset``, ``workspace_concept_fixture.py``): its parents, saved pairs, images, ratings and runs are what the Overview, the Statistics dialog and the health check read. These per-server hooks only change what a System test needs on top: - ``POST /__test__/system/reset`` restores the configured sources, forgets the session's health report and ends every injected failure; - ``POST /__test__/system/fail`` ``{path, times}`` answers the next ``times`` requests for ``path`` (a ``/prompt_manager/system/`` route) with the routes' 500 error envelope; - ``POST /__test__/system/missing-source`` configures one more source that does not exist (an unmounted drive); with ``{"inside": true}`` it is an existing folder inside the output folder (``output/unplugged``), which ``POST /__test__/system/unplug`` then removes (a drive unplugged while the page shows it); - ``POST /__test__/system/missing-file`` deletes the newest linked image's file from disk, leaving its database row (a missing-file link); - ``POST /__test__/system/unthumbnailed`` ``{count}`` writes ``count`` plain PNGs with no thumbnail into the output folder; - ``POST /__test__/system/monitor`` ``{watching, alive}`` stands in a fake image monitor whose observer is ``alive`` (default true) and schedules the named folders of this server's root (``["output"]``), so Sources reads watching, manual scan and restart pending states (S4). - ``POST /__test__/system/library`` ``{state}`` (S6b): ``node_folder`` makes the live database read as a node-folder library that can move to ``user/__promptmanager``; Move library then runs a stand-in move that touches no file and ends with the library in the default folder. ``failing`` is the same library whose move fails verification. Reset restores the real resolver, the real move and "Keep it here". - ``POST /__test__/system/legacy-previews`` ``{count}`` (#317) writes ``count`` PNGs into ``output/legacy`` with an older release's preview beside each (``thumbnails/legacy/_thumb.png``, newer than its file): outdated previews for Convert outdated; - ``POST /__test__/system/jxl-encoder`` ``{available}`` (#317) makes the server read as having (``cjxl``) or lacking a JPEG XL encoder without probing; reset restores the thumbnail settings and the real detection. - ``POST /__test__/system/duplicate-prompts`` (S7) adds two prompts whose text differs only in case and spacing (a duplicate prompt text group); ``{"again": true}`` adds a third, so a reviewed plan no longer matches; - ``POST /__test__/system/duplicate-files`` (S7) writes one PNG's bytes three times: ``output/s7-dupes/same.png``, ``input/s7-dupes/same.png`` (the same relative path in the other source) and ``output/s7-dupes/copy.png``; ``{"groups": n}`` adds ``n`` more groups of two copies each, and ``{"trio": false}`` leaves the three ``same.png``/``copy.png`` copies out, and ``{"long": true}`` adds a group under long names with no spaces; - ``GET /__test__/system/digest`` (S7) answers the sha256 of the database's rows (a read-only dump) and of every media file under ``output`` and ``input`` (not the automatic previews in ``thumbnails``), so a test proves a Cancel wrote nothing. - ``POST /__test__/system/safety-backup`` (S9) writes one verified safety backup beside the database, as a maintenance apply does; ``GET /__test__/system/backups`` answers ``{safety: [names], stages: n}``: the ``.backup_*`` files and the held ``.pm-restore-*`` stage folders. Reset drops every held restore review, removes the safety backups and forgets the session's last download. Delays use the shared request holds (``request_hold.py``). """ import hashlib import os import shutil import sqlite3 from pathlib import Path from aiohttp import web from PIL import Image FAILABLE_PREFIX = "/prompt_manager/system/" MAX_UNTHUMBNAILED = 50 def register_system_commands(app, routes, api, root: Path): # System › Appearance saves the host fields to this server's config.json (S3); # reset restores both. The file is this server's own, inside its temp root: # reset rewrites or deletes it, so a server whose configuration points # anywhere else (a real config.json) refuses to register these hooks. config_path = root / "config.json" configured_path = os.environ.get("PROMPT_MANAGER_CONFIG_PATH") if ( configured_path is None or Path(configured_path).resolve() != config_path.resolve() ): raise RuntimeError( f"PROMPT_MANAGER_CONFIG_PATH must be {config_path}, not {configured_path!r}" ) from py.config import GalleryConfig, PromptManagerConfig from utils import image_monitor from database import library_move from utils import library_location configured = list(GalleryConfig.MONITORING_DIRECTORIES) real_move = library_move.move_library keep = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER config_bytes = config_path.read_bytes() if config_path.exists() else None host_values = ( PromptManagerConfig.RESULT_TIMEOUT, PromptManagerConfig.WEBUI_DISPLAY_MODE, PromptManagerConfig.WORKER_THREADS, GalleryConfig.MONITORING_ENABLED, PromptManagerConfig.THUMBNAIL_FORMAT, PromptManagerConfig.THUMBNAIL_QUALITY, PromptManagerConfig.THUMBNAIL_EDGE, ) from py.api import thumbnail_encode monitor = getattr(image_monitor, "_monitor_instance", None) failures: dict[str, int] = {} @web.middleware async def inject_failures(request, handler): left = failures.get(request.path, 0) if left <= 0: return await handler(request) failures[request.path] = left - 1 return web.json_response( {"success": False, "code": None, "error": "Injected failure"}, status=500 ) app.middlewares.append(inject_failures) unplugged = root / "output" / "unplugged" dupes = [root / "output" / "s7-dupes", root / "input" / "s7-dupes"] def reset(): failures.clear() for folder in dupes: shutil.rmtree(folder, ignore_errors=True) if unplugged.is_dir(): unplugged.rmdir() GalleryConfig.MONITORING_DIRECTORIES = list(configured) api._health_report = None ( PromptManagerConfig.RESULT_TIMEOUT, PromptManagerConfig.WEBUI_DISPLAY_MODE, PromptManagerConfig.WORKER_THREADS, GalleryConfig.MONITORING_ENABLED, PromptManagerConfig.THUMBNAIL_FORMAT, PromptManagerConfig.THUMBNAIL_QUALITY, PromptManagerConfig.THUMBNAIL_EDGE, ) = host_values thumbnail_encode.reset_encoders() image_monitor._monitor_instance = monitor api.__dict__.pop("_library_resolver", None) library_move.move_library = real_move PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep if config_bytes is None: config_path.unlink(missing_ok=True) else: # Replace in one step, as the server writes it (Windows preflight). staged = config_path.with_name(config_path.name + ".reset") staged.write_bytes(config_bytes) os.replace(staged, config_path) def backup_files(): db_path = Path(api.db.model.db_path) safety = sorted( path.name for path in db_path.parent.glob(db_path.name + ".backup_*") if path.is_file() ) stages = [ path for path in db_path.parent.glob(".pm-restore-*") if path.is_dir() and not path.name.startswith(".pm-restore-upload-") ] return db_path, safety, stages async def reset_backups(): await api._discard_held(api._restore_plans.pop_all()) db_path, safety, _stages = backup_files() for name in safety: (db_path.parent / name).unlink(missing_ok=True) api._last_backup_download = None @routes.post("/__test__/system/reset") async def system_reset(request): reset() await reset_backups() return web.json_response({"success": True}) @routes.post("/__test__/system/safety-backup") async def safety_backup(request): from py.api.maintenance import take_safety_backup name = await api._run_in_executor(take_safety_backup, api.db.model) return web.json_response({"success": True, "name": name}) @routes.get("/__test__/system/backups") async def backups(request): _db_path, safety, stages = backup_files() return web.json_response({"safety": safety, "stages": len(stages)}) @routes.post("/__test__/system/fail") async def system_fail(request): data = await request.json() path, times = data.get("path"), data.get("times", 1) if ( not isinstance(path, str) or not path.startswith(FAILABLE_PREFIX) or isinstance(times, bool) or not isinstance(times, int) or not 1 <= times <= 10 ): raise web.HTTPBadRequest() failures[path] = times return web.json_response({"success": True}) @routes.post("/__test__/system/missing-source") async def missing_source(request): data = await request.json() if request.can_read_body else {} if data.get("inside") is True: unplugged.mkdir(exist_ok=True) absent = str(unplugged) else: absent = str(root / "unmounted-drive") if os.path.exists(absent): raise web.HTTPConflict() GalleryConfig.MONITORING_DIRECTORIES = [*configured, absent] return web.json_response( {"success": True, "configured": len(GalleryConfig.MONITORING_DIRECTORIES)} ) @routes.post("/__test__/system/unplug") async def unplug(request): if not unplugged.is_dir(): raise web.HTTPConflict() unplugged.rmdir() return web.json_response({"success": True}) @routes.post("/__test__/system/missing-file") async def missing_file(request): with api.db.model.get_connection() as conn: row = conn.execute( "SELECT file_path FROM generated_images ORDER BY id DESC LIMIT 1" ).fetchone() if row is None: raise web.HTTPNotFound() Path(row[0]).unlink() api.invalidate_gallery_cache() return web.json_response({"success": True, "filename": Path(row[0]).name}) @routes.post("/__test__/system/unthumbnailed") async def unthumbnailed(request): data = await request.json() count = data.get("count") if ( isinstance(count, bool) or not isinstance(count, int) or not 1 <= count <= MAX_UNTHUMBNAILED ): raise web.HTTPBadRequest() folder = root / "output" / "unthumbnailed" folder.mkdir(exist_ok=True) for number in range(count): with Image.new("RGB", (64, 64), "#2b3442") as image: image.save(folder / f"plain_{number + 1:03}.png") api.invalidate_gallery_cache() return web.json_response({"success": True}) @routes.post("/__test__/system/legacy-previews") async def legacy_previews(request): data = await request.json() count = data.get("count") if ( isinstance(count, bool) or not isinstance(count, int) or not 1 <= count <= MAX_UNTHUMBNAILED ): raise web.HTTPBadRequest() folder = root / "output" / "legacy" previews = root / "output" / "thumbnails" / "legacy" folder.mkdir(exist_ok=True) previews.mkdir(parents=True, exist_ok=True) names = [] for number in range(count): name = f"old_{number + 1:03}" with Image.effect_noise((256, 256), 50) as noise: noise.convert("RGB").save(folder / f"{name}.png") noise.convert("RGB").save(previews / f"{name}_thumb.png") later = (folder / f"{name}.png").stat().st_mtime + 5 os.utime(previews / f"{name}_thumb.png", (later, later)) names.append(f"legacy/{name}.png") api.invalidate_gallery_cache() return web.json_response({"success": True, "files": names}) @routes.post("/__test__/system/jxl-encoder") async def jxl_encoder(request): data = await request.json() available = data.get("available") if not isinstance(available, bool): raise web.HTTPBadRequest() # Test-only: stand in the cached detection (no probe). with thumbnail_encode._detect_lock: thumbnail_encode._detected = "cjxl" if available else None return web.json_response({"success": True}) @routes.post("/__test__/system/monitor") async def fake_monitor(request): data = await request.json() watching, alive = data.get("watching", []), data.get("alive", True) if ( not isinstance(watching, list) or not all(name in ("output", "input") for name in watching) or not isinstance(alive, bool) ): raise web.HTTPBadRequest() image_monitor._monitor_instance = _FakeMonitor( [str(root / name) for name in watching], alive ) return web.json_response({"success": True}) @routes.post("/__test__/system/duplicate-prompts") async def duplicate_prompts(request): data = await request.json() if request.can_read_body else {} texts = [("S7 duplicate prompt", "2001-01-01 00:00:00")] texts.append((" s7 DUPLICATE prompt", "2002-01-01 00:00:00")) if data.get("again") is True: texts = [("S7 Duplicate Prompt ", "2003-01-01 00:00:00")] with api.db.model.get_connection() as conn: conn.executemany( "INSERT INTO prompts(text, created_at) VALUES(?, ?)", texts ) conn.commit() return web.json_response({"success": True}) @routes.post("/__test__/system/duplicate-files") async def duplicate_files(request): body = await request.json() if request.can_read_body else {} for folder in dupes: folder.mkdir(exist_ok=True) extra = body.get("groups", 0) if ( isinstance(extra, bool) or not isinstance(extra, int) or not 0 <= extra <= 20 ): raise web.HTTPBadRequest() for number in range(extra): # One more group each: two copies of a picture no other file has. with Image.new("RGB", (16, 16), (number * 9, 40, 90)) as image: image.save(dupes[0] / f"group_{number:02}_a.png") copy = (dupes[0] / f"group_{number:02}_a.png").read_bytes() (dupes[0] / f"group_{number:02}_b.png").write_bytes(copy) if body.get("long") is True: # One group whose copies sit under long folder and file names with # no space to break at, as real dataset folders do. deep = dupes[0] / ("mabes9sans_DATASET_ideogram_v1_" * 3).rstrip("_") deep.mkdir(exist_ok=True) with Image.new("RGB", (16, 16), "#204060") as image: image.save(deep / ("mabes9sans_DATASET072_" * 4 + "a.png")) copy = (deep / ("mabes9sans_DATASET072_" * 4 + "a.png")).read_bytes() (deep / ("mabes9sans_DATASET072_" * 4 + "b.png")).write_bytes(copy) if body.get("trio", True) is not False: with Image.new("RGB", (48, 48), "#7a3b2e") as image: image.save(dupes[0] / "same.png") data = (dupes[0] / "same.png").read_bytes() (dupes[1] / "same.png").write_bytes(data) (dupes[0] / "copy.png").write_bytes(data) api.invalidate_gallery_cache() return web.json_response({"success": True}) @routes.get("/__test__/system/digest") async def digest(request): uri = Path(api.db.model.db_path).as_uri() + "?mode=ro" rows = hashlib.sha256() conn = sqlite3.connect(uri, uri=True) try: for line in conn.iterdump(): rows.update(line.encode("utf-8") + b"\n") finally: conn.close() media = hashlib.sha256() for name in ("output", "input"): for path in sorted((root / name).rglob("*")): # Automatic previews are written in the background (#316): media only. if path.is_file() and "thumbnails" not in path.relative_to(root).parts: media.update(str(path.relative_to(root)).encode("utf-8")) media.update(hashlib.sha256(path.read_bytes()).digest()) return web.json_response({"rows": rows.hexdigest(), "media": media.hexdigest()}) @routes.post("/__test__/system/library") async def library_state(request): data = await request.json() state = data.get("state") if state not in ("node_folder", "failing"): raise web.HTTPBadRequest() target = str(root / "user" / "__promptmanager") def located(source): return library_location.LibraryLocation( source, api.db.model.db_path, str(config_path), str(root), target if source == library_location.SOURCE_NODE_FOLDER else None, target_source=library_location.SOURCE_DEFAULT, ) current = {"source": library_location.SOURCE_NODE_FOLDER} api._library_resolver = lambda: located(current["source"]) def stand_in_move(plan, *, progress, **_kwargs): if state == "failing": raise library_move.MoveFailed("verify-failed") progress("Copying the database") counts = api.db.library_counts() current["source"] = library_location.SOURCE_DEFAULT return library_move.MoveResult( located(library_location.SOURCE_DEFAULT), {"prompts": counts["positive_prompts"], "generated_images": 0}, ["prompts.db"], ["browser.db.moved-to-__promptmanager"], ) library_move.move_library = stand_in_move return web.json_response({"success": True}) return reset class _FakeMonitor: """The image monitor's status surface (``get_status``) with no observer thread.""" def __init__(self, directories, alive): self.monitored_directories = directories self._alive = alive def get_status(self): return { "running": True, "monitored_directories": list(self.monitored_directories), "handler_active": True, "observer_alive": self._alive, }