480 lines
19 KiB
Python
480 lines
19 KiB
Python
"""System routes (py/api/system.py): health check job, Overview and statistics.
|
|
|
|
The health check runs off-loop as a ``read`` registry job of kind ``health``
|
|
(one at a time, beside a scan), and the API keeps its last report. The
|
|
Overview reads only real values: Library-eligible counts, indexed images, the
|
|
last health result or null, no backup until one exists, and the effective
|
|
sources and whether they were auto-detected. The missing-thumbnail count is not
|
|
part of it: the Overview reads ``/images/thumbnails/status`` on its own (its
|
|
walk of every source must not hold the other values back).
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import threading
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
_mock_server = MagicMock()
|
|
_mock_server.PromptServer.instance.routes = MagicMock()
|
|
sys.modules.setdefault("server", _mock_server)
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from py.api.jobs import HEALTH_JOB_KIND # noqa: E402
|
|
from py.config import GalleryConfig # noqa: E402
|
|
from utils.hashing import generate_prompt_hash # noqa: E402
|
|
|
|
try:
|
|
from tests.test_api_admin import AdminAPITestCase # noqa: E402
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_api_admin import AdminAPITestCase # noqa: E402
|
|
|
|
WAIT = 5
|
|
|
|
|
|
class SystemAPITestCase(AdminAPITestCase):
|
|
async def get_application(self):
|
|
app = await super().get_application()
|
|
# Auto-detection searches the disk; here it finds the fake ComfyUI only.
|
|
self.api._find_comfyui_output_dir = lambda: (
|
|
str(self.output_dir) if self.output_dir.is_dir() else None
|
|
)
|
|
return app
|
|
|
|
def save(self, text, **metadata):
|
|
return self.api.db.save_prompt(
|
|
text, prompt_hash=generate_prompt_hash(text), **metadata
|
|
)
|
|
|
|
def image(self, prompt_id, name, *, exists=True, folder=None):
|
|
path = Path(folder or self.output_dir) / name
|
|
if exists:
|
|
path.write_bytes(b"png")
|
|
return self.api.db.link_image_to_prompt(prompt_id, str(path))
|
|
|
|
async def json(self, method, url, status=200):
|
|
response = await self.client.request(method, url)
|
|
self.assertEqual(response.status, status, await response.text())
|
|
return await response.json()
|
|
|
|
async def restore_snapshot(self):
|
|
"""Restore the live database from a copy of itself through the legacy route."""
|
|
from contextlib import closing
|
|
|
|
from aiohttp import FormData
|
|
|
|
snapshot = os.path.join(self.tmpdir, "snapshot.db")
|
|
with closing(sqlite3.connect(self.api.db.model.db_path)) as src:
|
|
with closing(sqlite3.connect(snapshot)) as dst:
|
|
src.backup(dst)
|
|
with open(snapshot, "rb") as handle:
|
|
upload = handle.read()
|
|
form = FormData()
|
|
form.add_field("database_file", upload, filename="upload.db")
|
|
response = await self.client.request(
|
|
"POST", "/prompt_manager/restore", data=form
|
|
)
|
|
self.assertEqual(response.status, 200, await response.text())
|
|
|
|
async def run_health(self):
|
|
body = await self.json("POST", "/prompt_manager/system/health", 202)
|
|
job = self.api._jobs.get(body["job"]["id"])
|
|
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
|
return body["job"], job
|
|
|
|
|
|
class TestHealthRoute(SystemAPITestCase):
|
|
async def test_no_report_before_a_check(self):
|
|
body = await self.json("GET", "/prompt_manager/system/health")
|
|
self.assertEqual(body, {"success": True, "report": None})
|
|
|
|
async def test_health_check_runs_as_a_read_job_and_keeps_its_report(self):
|
|
prompt = self.save("Images")
|
|
self.image(prompt, "kept.png")
|
|
self.image(prompt, "gone.png", exists=False)
|
|
summary, job = await self.run_health()
|
|
self.assertEqual(summary["kind"], HEALTH_JOB_KIND)
|
|
self.assertEqual(summary["title"], "Health check")
|
|
self.assertEqual(summary["admission"], "read")
|
|
self.assertFalse(summary["cancellable"])
|
|
self.assertEqual(summary["counters"], {})
|
|
self.assertEqual(job.state, "finished")
|
|
self.assertEqual(job.progress, 100)
|
|
self.assertEqual(job.status_line, "Health check incomplete")
|
|
report = (await self.json("GET", "/prompt_manager/system/health"))["report"]
|
|
self.assertEqual(report["status"], "incomplete")
|
|
missing = next(c for c in report["checks"] if c["id"] == "missing_files")
|
|
self.assertEqual((missing["status"], missing["count"]), ("issues", 1))
|
|
ambiguous = next(c for c in report["checks"] if c["id"] == "ambiguous_metadata")
|
|
self.assertEqual(ambiguous["status"], "unavailable")
|
|
self.assertNotIn(self.tmpdir, str(report))
|
|
|
|
async def test_health_check_is_single_flight_and_runs_beside_a_scan(self):
|
|
release = threading.Event()
|
|
entered = threading.Event()
|
|
real = self.api.db.run_health_checks
|
|
|
|
def gated(*args, **kwargs):
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
return real(*args, **kwargs)
|
|
|
|
scan_gate = asyncio.Event()
|
|
|
|
async def scanning(job):
|
|
await scan_gate.wait()
|
|
|
|
scan = self.api._jobs.start(
|
|
"scan", "Metadata scan", admission="records", runner=scanning
|
|
)
|
|
with patch.object(self.api.db, "run_health_checks", side_effect=gated):
|
|
first = await self.json("POST", "/prompt_manager/system/health", 202)
|
|
loop = asyncio.get_running_loop()
|
|
await loop.run_in_executor(None, entered.wait, WAIT)
|
|
busy = await self.json("POST", "/prompt_manager/system/health", 409)
|
|
self.assertEqual(busy["code"], "busy")
|
|
self.assertEqual(
|
|
busy["running"], {"id": first["job"]["id"], "title": "Health check"}
|
|
)
|
|
release.set()
|
|
await asyncio.wait_for(self.api._jobs.get(first["job"]["id"]).task, WAIT)
|
|
scan_gate.set()
|
|
await asyncio.wait_for(scan.task, WAIT)
|
|
|
|
async def test_failed_check_keeps_the_previous_report(self):
|
|
await self.run_health()
|
|
before = (await self.json("GET", "/prompt_manager/system/health"))["report"]
|
|
with (
|
|
patch.object(
|
|
self.api.db, "run_health_checks", side_effect=RuntimeError("boom")
|
|
),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
body = await self.json("POST", "/prompt_manager/system/health", 202)
|
|
job = self.api._jobs.get(body["job"]["id"])
|
|
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
|
self.assertEqual(job.state, "failed")
|
|
after = (await self.json("GET", "/prompt_manager/system/health"))["report"]
|
|
self.assertEqual(after, before)
|
|
|
|
async def test_absent_configured_source_is_reported_per_source(self):
|
|
prompt = self.save("Images")
|
|
self.image(prompt, "gone.png", exists=False)
|
|
absent = self.comfy_dir / "output" / "unmounted"
|
|
GalleryConfig.MONITORING_DIRECTORIES = [
|
|
self._root(self.output_dir),
|
|
str(absent),
|
|
]
|
|
await self.run_health()
|
|
report = (await self.json("GET", "/prompt_manager/system/health"))["report"]
|
|
missing = next(c for c in report["checks"] if c["id"] == "missing_files")
|
|
# Per source: the readable output is counted, the absent one is not.
|
|
self.assertEqual(missing["status"], "issues")
|
|
self.assertEqual(missing["count"], 1)
|
|
self.assertTrue(missing["partial"])
|
|
self.assertEqual(
|
|
[(s["status"], s["count"]) for s in missing["sources"]],
|
|
[("checked", 1), ("unavailable", None)],
|
|
)
|
|
self.assertEqual(report["status"], "incomplete")
|
|
|
|
async def test_restore_clears_health_report(self):
|
|
self.save("Any")
|
|
await self.run_health()
|
|
self.assertIsNotNone(
|
|
(await self.json("GET", "/prompt_manager/system/health"))["report"]
|
|
)
|
|
await self.restore_snapshot()
|
|
body = await self.json("GET", "/prompt_manager/system/health")
|
|
self.assertEqual(body, {"success": True, "report": None})
|
|
overview = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertIsNone(overview["health"])
|
|
|
|
async def test_check_finishing_after_a_restore_is_not_kept(self):
|
|
# A check that started on the old database describes it, not the new one.
|
|
self.save("Any")
|
|
release = threading.Event()
|
|
entered = threading.Event()
|
|
real = self.api.db.run_health_checks
|
|
|
|
def gated(*args, **kwargs):
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
return real(*args, **kwargs)
|
|
|
|
with patch.object(self.api.db, "run_health_checks", side_effect=gated):
|
|
body = await self.json("POST", "/prompt_manager/system/health", 202)
|
|
loop = asyncio.get_running_loop()
|
|
await loop.run_in_executor(None, entered.wait, WAIT)
|
|
await self.restore_snapshot()
|
|
release.set()
|
|
job = self.api._jobs.get(body["job"]["id"])
|
|
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
|
self.assertEqual(job.state, "finished")
|
|
body = await self.json("GET", "/prompt_manager/system/health")
|
|
self.assertEqual(body, {"success": True, "report": None})
|
|
# A check started on the restored database is kept again.
|
|
await self.run_health()
|
|
self.assertIsNotNone(
|
|
(await self.json("GET", "/prompt_manager/system/health"))["report"]
|
|
)
|
|
|
|
async def test_health_sources_resolve_off_the_event_loop(self):
|
|
loop_thread = threading.get_ident()
|
|
threads = []
|
|
real = self.api._health_roots
|
|
|
|
def recording():
|
|
threads.append(threading.get_ident())
|
|
return real()
|
|
|
|
with patch.object(self.api, "_health_roots", side_effect=recording):
|
|
await self.run_health()
|
|
self.assertEqual(len(threads), 1)
|
|
self.assertNotEqual(threads[0], loop_thread)
|
|
|
|
async def test_health_start_failure_is_an_error_envelope(self):
|
|
with (
|
|
patch.object(self.api, "_health_roots", side_effect=RuntimeError("secret")),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
body = await self.json("POST", "/prompt_manager/system/health", 500)
|
|
self.assertEqual(
|
|
body,
|
|
{
|
|
"success": False,
|
|
"code": None,
|
|
"error": "Failed to start the health check",
|
|
},
|
|
)
|
|
self.assertEqual(self.api._jobs.list(), [])
|
|
|
|
|
|
class TestOverviewRoute(SystemAPITestCase):
|
|
async def test_overview_reports_real_counts(self):
|
|
self.api.db.save_prompt_pair("Cat", "Blur")
|
|
archived = self.api.db.save_prompt_pair("Cat", "Noise")
|
|
self.api.db.archive_prompt_pair(archived)
|
|
legacy = self.save("Legacy")
|
|
self.image(legacy, "a.png")
|
|
self.image(legacy, "b.png")
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertTrue(body["success"])
|
|
self.assertEqual(
|
|
body["library"], {"positive_prompts": 2, "negative_variants": 1}
|
|
)
|
|
self.assertEqual(body["images"], {"indexed": 2, "rated": 0})
|
|
self.assertEqual(
|
|
set(body),
|
|
{
|
|
"success",
|
|
"library",
|
|
"images",
|
|
"health",
|
|
"backup",
|
|
"needs_review",
|
|
"sources",
|
|
},
|
|
)
|
|
|
|
async def test_overview_rated_counts_images_of_rated_prompts(self):
|
|
rated = self.save("Rated", rating=5)
|
|
self.image(rated, "a.png")
|
|
self.image(rated, "b.png")
|
|
self.image(self.save("Unrated"), "c.png")
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertEqual(body["images"], {"indexed": 3, "rated": 2})
|
|
|
|
async def test_overview_never_fabricates_backup_or_health(self):
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertIsNone(body["health"])
|
|
self.assertIsNone(body["backup"])
|
|
self.assertEqual(body["needs_review"], {"pending": 0})
|
|
self.save("Any")
|
|
await self.run_health()
|
|
report = (await self.json("GET", "/prompt_manager/system/health"))["report"]
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertEqual(
|
|
body["health"],
|
|
{"checked_at": report["checked_at"], "status": "incomplete", "issues": 0},
|
|
)
|
|
self.assertIsNone(body["backup"])
|
|
|
|
async def test_overview_health_counts_checks_with_issues(self):
|
|
prompt = self.save("Images")
|
|
self.image(prompt, "gone.png", exists=False)
|
|
self.image(prompt, "gone-too.png", exists=False)
|
|
await self.run_health()
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertEqual(body["health"]["issues"], 1)
|
|
|
|
async def test_overview_sources_distinguish_auto_detect(self):
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertEqual(
|
|
body["sources"], {"configured": 0, "effective": 1, "auto_detected": True}
|
|
)
|
|
second = self.output_dir / "archive"
|
|
second.mkdir()
|
|
GalleryConfig.MONITORING_DIRECTORIES = [
|
|
self._root(self.output_dir),
|
|
self._root(second),
|
|
]
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertEqual(
|
|
body["sources"], {"configured": 2, "effective": 2, "auto_detected": False}
|
|
)
|
|
# Every configured source unavailable: ComfyUI's output is used instead.
|
|
GalleryConfig.MONITORING_DIRECTORIES = [str(self.comfy_dir / "absent")]
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertEqual(
|
|
body["sources"], {"configured": 1, "effective": 1, "auto_detected": True}
|
|
)
|
|
|
|
async def test_overview_never_walks_the_sources_for_thumbnails(self):
|
|
(self.output_dir / "one.png").write_bytes(b"png")
|
|
with patch.object(
|
|
self.api,
|
|
"_plan_gallery_thumbnails",
|
|
side_effect=AssertionError("walked"),
|
|
):
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertTrue(body["success"])
|
|
self.assertNotIn("thumbnails", body)
|
|
|
|
async def test_overview_sources_resolve_off_the_event_loop(self):
|
|
loop_thread = threading.get_ident()
|
|
threads = []
|
|
real = self.api._output_dirs_with_origin
|
|
|
|
def recording():
|
|
threads.append(threading.get_ident())
|
|
return real()
|
|
|
|
with patch.object(self.api, "_output_dirs_with_origin", side_effect=recording):
|
|
body = await self.json("GET", "/prompt_manager/system/overview")
|
|
self.assertEqual(body["sources"]["effective"], 1)
|
|
self.assertEqual(len(threads), 1)
|
|
self.assertNotEqual(threads[0], loop_thread)
|
|
|
|
async def test_overview_failure_is_an_error_envelope(self):
|
|
with (
|
|
patch.object(
|
|
self.api.db, "library_counts", side_effect=RuntimeError("secret")
|
|
),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
body = await self.json("GET", "/prompt_manager/system/overview", 500)
|
|
self.assertEqual(
|
|
body,
|
|
{"success": False, "code": None, "error": "Failed to read the overview"},
|
|
)
|
|
|
|
|
|
class TestStatisticsRoute(SystemAPITestCase):
|
|
async def test_statistics_sum_real_run_counts(self):
|
|
a = self.save("Alpha", category="Landscapes", rating=5)
|
|
b = self.save("Beta", category="Portraits", rating=2)
|
|
self.api.db.record_prompt_use(a, times=3)
|
|
self.api.db.record_prompt_use(b)
|
|
body = await self.json("GET", "/prompt_manager/system/statistics")
|
|
self.assertEqual(
|
|
body["by_category"],
|
|
[
|
|
{"category": "Landscapes", "runs": 3, "average_rating": 5.0},
|
|
{"category": "Portraits", "runs": 1, "average_rating": 2.0},
|
|
],
|
|
)
|
|
self.assertEqual(body["totals"], self.api.db.get_statistics())
|
|
self.assertTrue(body["success"])
|
|
|
|
async def test_statistics_lists_rated_images_for_library_totals(self):
|
|
# D96: the Statistics dialog's "Library totals" reads the rated count
|
|
# with the runs (one read), from the same images the Overview counts.
|
|
rated = self.save("Rated", rating=4)
|
|
self.image(rated, "a.png")
|
|
self.image(self.save("Unrated"), "b.png")
|
|
body = await self.json("GET", "/prompt_manager/system/statistics")
|
|
self.assertEqual(body["rated_images"], 1)
|
|
self.assertEqual(
|
|
set(body), {"success", "by_category", "totals", "rated_images"}
|
|
)
|
|
|
|
async def test_statistics_failure_is_an_error_envelope(self):
|
|
with (
|
|
patch.object(
|
|
self.api.db, "usage_by_category", side_effect=RuntimeError("secret")
|
|
),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
body = await self.json("GET", "/prompt_manager/system/statistics", 500)
|
|
self.assertEqual(
|
|
body,
|
|
{"success": False, "code": None, "error": "Failed to read statistics"},
|
|
)
|
|
|
|
|
|
class TestFailedReadsSettleFirst(SystemAPITestCase):
|
|
"""A failed read answers only after the route's other reads finished.
|
|
|
|
Answering on the first failure (``asyncio.gather`` without
|
|
``return_exceptions``) left the other executor reads running against a
|
|
database the caller may already be closing; on Python 3.10 CI that
|
|
crashed the test process inside SQLite.
|
|
"""
|
|
|
|
async def assert_answers_after_reads(self, url, failing, slow, error):
|
|
raised, release, finished = (threading.Event() for _ in range(3))
|
|
real = getattr(*slow)
|
|
|
|
def fail(*args, **kwargs):
|
|
raised.set()
|
|
raise RuntimeError("secret")
|
|
|
|
def held(*args, **kwargs):
|
|
release.wait()
|
|
try:
|
|
return real(*args, **kwargs)
|
|
finally:
|
|
finished.set()
|
|
|
|
with (
|
|
patch.object(*failing, side_effect=fail),
|
|
patch.object(*slow, side_effect=held),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
answer = asyncio.ensure_future(self.json("GET", url, 500))
|
|
try:
|
|
await asyncio.get_running_loop().run_in_executor(None, raised.wait)
|
|
done, _ = await asyncio.wait({answer}, timeout=0.5)
|
|
self.assertEqual(done, set(), "answered while a read still ran")
|
|
finally:
|
|
release.set()
|
|
body = await answer
|
|
self.assertTrue(finished.is_set())
|
|
self.assertEqual(body, {"success": False, "code": None, "error": error})
|
|
|
|
async def test_statistics(self):
|
|
await self.assert_answers_after_reads(
|
|
"/prompt_manager/system/statistics",
|
|
(self.api.db, "usage_by_category"),
|
|
(self.api.db, "get_statistics"),
|
|
"Failed to read statistics",
|
|
)
|
|
|
|
async def test_overview(self):
|
|
await self.assert_answers_after_reads(
|
|
"/prompt_manager/system/overview",
|
|
(self.api.db, "library_counts"),
|
|
(self.api, "_output_dirs_with_origin"),
|
|
"Failed to read the overview",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import unittest
|
|
|
|
unittest.main()
|