Vito Sansevero
ea23d5161d
D129. System > Overview's Library health gains a "Library location" row (the database path, why it is there, the "Keep it here" choice and any PROMPT_MANAGER_DATA_DIR problem) with Move library for a node-folder library and Details otherwise. While a node-folder library is not kept there, the workspace shows the concept's .notice above every page: "Keep it here" saves the choice on the server and stops the notice; "Move library" opens the shared review (From, To, Size, what moves, what happens to the old files) and starts the exclusive job in Activity. The notice and the Overview read the location again when the job ends. The Overview comparison maps the inserted row; a new comparison holds the notice to the concept's Sources notice (markup, tokens, alignment, both themes and widths) with a mutation check.
258 lines
10 KiB
Python
258 lines
10 KiB
Python
"""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".
|
||
|
||
Delays use the shared request holds (``request_hold.py``).
|
||
"""
|
||
|
||
import os
|
||
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,
|
||
)
|
||
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"
|
||
|
||
def reset():
|
||
failures.clear()
|
||
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,
|
||
) = host_values
|
||
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)
|
||
|
||
@routes.post("/__test__/system/reset")
|
||
async def system_reset(request):
|
||
reset()
|
||
return web.json_response({"success": True})
|
||
|
||
@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/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/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,
|
||
}
|