Vito
b1e7abd380
* fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370) - PromptTracker: the cleanup thread waits on a stop event instead of time.sleep; new shutdown(timeout=2) joins it and returns whether it ended. - ImageGenerationHandler.shutdown(timeout=5) now joins the worker and returns a bool; new close() is final: schedule() refuses and a sealed run's history re-check is dropped, so no worker restarts afterwards. - ImageMonitor.stop_monitoring(timeout=5) closes the handler first, then stops and joins the watchdog observer within one deadline. - Every test that builds a real tracker, handler or monitor ends with its final stop and an assertion, before its folder and databases go. - #370: the wrapper/session history test drains before dispose(); the held-write bound test uses 30 s ceilings instead of tight wall clocks. * test(history): drain the writer before asserting a bounded close it does not test (#370) * fix(threads): a later stop waits again for a worker or observer an earlier stop left running (#302) Review findings on #372: - ImageGenerationHandler keeps the workers a stop signalled that had not ended; every later shutdown/close joins them again, so a second call no longer returns True while the first worker still runs. - ImageMonitor.stop_monitoring keeps an observer or handler it left running and stops and joins it again on the next call. - A closed handler takes no awaiting history path, and a re-check that raced close() is not kept due. - Tests: second stops of a busy worker and a live observer, the closed handler's re-check paths, the handle test asserts the worker ended and stops before it asserts, and a no-new-thread check instead of set equality. * test(browser): the System reset restores every settings value, safety_backup included The frontend coverage shard of #372 failed in workspace-system-tagging: the review read "No safety backup (the policy is off)". The backups spec saves safety_backup=false, and /__test__/system/reset restored config.json but not PromptManagerConfig.SAFETY_BACKUP (nor INFINITE_SCROLL), so a later spec on the same worker server saw the policy off. The reset now captures and restores every non-tagging value POST /settings can change.
444 lines
19 KiB
Python
444 lines
19 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".
|
||
- ``POST /__test__/system/legacy-previews`` ``{count}`` (#317) writes ``count``
|
||
PNGs into ``output/legacy`` with an older release's preview beside each
|
||
(``thumbnails/legacy/<name>_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 ``<db>.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
|
||
from py.api.settings_routes import (
|
||
_AUTOTAG_SETTINGS,
|
||
_GALLERY_SETTINGS,
|
||
_SETTINGS_FILE_KEYS,
|
||
)
|
||
|
||
# Every value POST /settings can change, so a reset restores each one
|
||
# (a saved safety_backup=false once leaked into a later spec's review).
|
||
# The tagging settings are the tagging fixture's to reset.
|
||
host_values = {
|
||
(GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig, name): (
|
||
getattr(
|
||
GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig,
|
||
name,
|
||
)
|
||
)
|
||
for name in _SETTINGS_FILE_KEYS
|
||
if name not in _AUTOTAG_SETTINGS
|
||
}
|
||
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
|
||
for (owner, name), value in host_values.items():
|
||
setattr(owner, name, value)
|
||
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,
|
||
}
|