Vito Sansevero
f6df8a39b3
GET /system/diagnostics builds the report on the server, off the event loop and under a registry lease (409 busy while a restore runs): the project and Viewer.js versions, Python and the platform, the four legacy checks (one helper, legacy_checks, that GET /diagnostics also calls), public database info, source counts, the log configuration and statistics, the last health summary and the config.json-only keys from an explicit list. The CivitAI key is reported only as has_civitai_api_key; every path is relative to ComfyUI or reads "outside ComfyUI". A scrub replaces anything still absolute, the key or the home folder and counts it in redacted (0 for a correct build). The version parser moves unchanged to utils/version.py. (#329)
303 lines
11 KiB
Python
303 lines
11 KiB
Python
"""The shareable diagnostic report (#329, py/api/diagnostics_report.py).
|
|
|
|
``GET /prompt_manager/system/diagnostics`` builds the report on the server
|
|
from an explicit list of fields: the project version, the four legacy checks,
|
|
public paths only and no secret. A scrub replaces anything that still looks
|
|
like an absolute path or holds the CivitAI key and counts it in ``redacted``;
|
|
these tests require ``redacted == 0`` so a builder mistake fails instead of
|
|
being hidden.
|
|
"""
|
|
|
|
import json
|
|
import ntpath
|
|
import os
|
|
import posixpath
|
|
import re
|
|
import sqlite3
|
|
import sys
|
|
import unittest
|
|
from unittest import mock
|
|
from unittest.mock import MagicMock
|
|
|
|
_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 import admin, diagnostics_report # noqa: E402
|
|
from py.api.diagnostics_report import report_path, scrub # noqa: E402
|
|
from py.config import (
|
|
GalleryConfig,
|
|
IntegrationConfig,
|
|
PromptManagerConfig,
|
|
) # 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
|
|
|
|
REPORT = "/prompt_manager/system/diagnostics"
|
|
KEY = "civitai-secret-0123456789abcdef"
|
|
SEEDED_PATHS = (
|
|
"/home/someone/private/LoraManager",
|
|
"/srv/someone/renders",
|
|
"C:\\Users\\someone\\Documents\\renders",
|
|
"D:/Users/someone/prompts.db",
|
|
"\\\\nas\\share\\someone\\renders",
|
|
)
|
|
_WINDOWS_ABSOLUTE = re.compile(r"^[A-Za-z]:[\\/]")
|
|
TOP_KEYS = {
|
|
"success",
|
|
"generated_at",
|
|
"version",
|
|
"python",
|
|
"platform",
|
|
"viewerjs",
|
|
"checks",
|
|
"database_info",
|
|
"sources",
|
|
"logging",
|
|
"health",
|
|
"config",
|
|
"redacted",
|
|
}
|
|
CONFIG = {
|
|
"database": {"enable_duplicate_detection", "enable_auto_save"},
|
|
"web_ui": {
|
|
"result_timeout",
|
|
"show_test_button",
|
|
"webui_display_mode",
|
|
"infinite_scroll",
|
|
},
|
|
"performance": {
|
|
"max_search_results",
|
|
"enable_fuzzy_search",
|
|
"auto_backup_interval",
|
|
"worker_threads",
|
|
},
|
|
"gallery": {
|
|
"processing_delay",
|
|
"prompt_timeout",
|
|
"cleanup_interval",
|
|
"auto_cleanup_missing_files",
|
|
"max_image_age_days",
|
|
"max_concurrent_processing",
|
|
"metadata_extraction_timeout",
|
|
},
|
|
"integrations": {"lora_manager_enabled", "has_civitai_api_key"},
|
|
}
|
|
|
|
|
|
def _strings(value):
|
|
if isinstance(value, str):
|
|
yield value
|
|
elif isinstance(value, dict):
|
|
for key, item in value.items():
|
|
yield str(key)
|
|
yield from _strings(item)
|
|
elif isinstance(value, (list, tuple)):
|
|
for item in value:
|
|
yield from _strings(item)
|
|
|
|
|
|
class DiagnosticsReportTestCase(AdminAPITestCase):
|
|
async def get_application(self):
|
|
app = await super().get_application()
|
|
saved = IntegrationConfig.get_config()
|
|
self.addCleanup(IntegrationConfig.update_config, saved)
|
|
return app
|
|
|
|
async def report(self, status=200):
|
|
response = await self.client.request("GET", REPORT)
|
|
text = await response.text()
|
|
self.assertEqual(response.status, status, text)
|
|
return json.loads(text), text
|
|
|
|
def seed_secrets(self):
|
|
IntegrationConfig.update_config(
|
|
{
|
|
"lora_manager": {
|
|
"enabled": True,
|
|
"path": SEEDED_PATHS[0],
|
|
"civitai_api_key": KEY,
|
|
}
|
|
}
|
|
)
|
|
GalleryConfig.MONITORING_DIRECTORIES = list(SEEDED_PATHS[1:3]) + [
|
|
SEEDED_PATHS[4]
|
|
]
|
|
PromptManagerConfig.update_config(
|
|
{"database": {"default_path": SEEDED_PATHS[3]}}
|
|
)
|
|
|
|
def assert_no_paths_or_key(self, report, text):
|
|
self.assertNotIn(KEY, text)
|
|
self.assertNotIn(self.tmpdir, text)
|
|
self.assertNotIn(self.tmpdir.replace("\\", "/"), text)
|
|
for seeded in SEEDED_PATHS:
|
|
self.assertNotIn(seeded, text)
|
|
self.assertNotIn(json.dumps(seeded)[1:-1], text)
|
|
self.assertNotIn("someone", text)
|
|
for value in _strings(report):
|
|
self.assertFalse(posixpath.isabs(value), value)
|
|
self.assertFalse(ntpath.isabs(value), value)
|
|
self.assertIsNone(_WINDOWS_ABSOLUTE.match(value), value)
|
|
self.assertFalse(value.startswith("\\\\"), value)
|
|
|
|
|
|
class TestReport(DiagnosticsReportTestCase):
|
|
async def test_report_has_version_viewerjs_and_the_four_checks(self):
|
|
report, _text = await self.report()
|
|
self.assertEqual(set(report), TOP_KEYS)
|
|
self.assertTrue(report["success"])
|
|
with open(
|
|
os.path.join(os.path.dirname(os.path.dirname(__file__)), "pyproject.toml"),
|
|
encoding="utf-8",
|
|
) as handle:
|
|
version = re.search(r'version\s*=\s*["\']([^"\']+)["\']', handle.read())
|
|
self.assertEqual(report["version"], version.group(1))
|
|
self.assertRegex(report["viewerjs"], r"^\d+\.\d+\.\d+$")
|
|
self.assertRegex(report["python"], r"^3\.\d+\.\d+")
|
|
self.assertEqual(set(report["platform"]), {"system", "release"})
|
|
self.assertEqual(
|
|
set(report["checks"]),
|
|
{"database", "dependencies", "comfyui_output", "image_monitor"},
|
|
)
|
|
self.assertEqual(report["checks"]["database"]["status"], "ok")
|
|
self.assertEqual(
|
|
set(report["sources"]),
|
|
{"configured", "effective", "watching", "unavailable"},
|
|
)
|
|
self.assertEqual(set(report["logging"]), {"config", "stats"})
|
|
self.assertIsNone(report["health"])
|
|
self.assertEqual(report["redacted"], 0)
|
|
self.assertRegex(report["generated_at"], r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d$")
|
|
|
|
async def test_report_has_no_secret_or_absolute_path(self):
|
|
self.seed_secrets()
|
|
with mock.patch.object(
|
|
admin.sqlite3,
|
|
"connect",
|
|
side_effect=sqlite3.DatabaseError("file is not a database"),
|
|
):
|
|
report, text = await self.report()
|
|
self.assertEqual(report["checks"]["database"]["status"], "error")
|
|
self.assertTrue(report["config"]["integrations"]["has_civitai_api_key"])
|
|
self.assertTrue(report["config"]["integrations"]["lora_manager_enabled"])
|
|
self.assertEqual(report["sources"]["configured"], 3)
|
|
self.assert_no_paths_or_key(report, text)
|
|
self.assertEqual(report["redacted"], 0)
|
|
|
|
async def test_report_leaves_get_config_unchanged(self):
|
|
self.seed_secrets()
|
|
before = json.dumps(PromptManagerConfig.get_config(), sort_keys=True)
|
|
directories = GalleryConfig.MONITORING_DIRECTORIES
|
|
await self.report()
|
|
self.assertEqual(
|
|
json.dumps(PromptManagerConfig.get_config(), sort_keys=True), before
|
|
)
|
|
self.assertIs(GalleryConfig.MONITORING_DIRECTORIES, directories)
|
|
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, KEY)
|
|
|
|
def test_outside_path_reads_outside_comfyui(self):
|
|
inside = os.path.join(str(self.comfy_dir), "user", "__promptmanager", "logs")
|
|
self.assertEqual(report_path(inside), "user/__promptmanager/logs")
|
|
self.assertEqual(
|
|
report_path(os.path.join(str(self.outside_dir), "vito")), "outside ComfyUI"
|
|
)
|
|
self.assertEqual(report_path(""), "")
|
|
self.assertEqual(report_path(None), "")
|
|
|
|
def test_scrub_removes_a_leaked_path_and_counts_it(self):
|
|
leaked = {
|
|
"a": "/home/someone/x",
|
|
"b": ["C:\\Users\\someone", "D:/data", "\\\\nas\\share", "ok"],
|
|
"c": {"d": f"token {KEY} here", "e": 3, "f": None, "g": "relative/path"},
|
|
"redacted": 0,
|
|
}
|
|
cleaned = scrub(leaked, KEY)
|
|
self.assertEqual(cleaned["redacted"], 5)
|
|
self.assertEqual(cleaned["a"], "[removed]")
|
|
self.assertEqual(cleaned["b"], ["[removed]"] * 3 + ["ok"])
|
|
self.assertEqual(
|
|
cleaned["c"], {"d": "[removed]", "e": 3, "f": None, "g": "relative/path"}
|
|
)
|
|
# A short key is not searched for (it would match ordinary words).
|
|
self.assertEqual(scrub({"x": "abc", "redacted": 0}, "abc")["redacted"], 0)
|
|
|
|
async def test_report_refused_while_exclusive_job_runs(self):
|
|
async with self.api._jobs.hold(
|
|
"restore", "Database restore", admission="exclusive"
|
|
) as restore:
|
|
report, _text = await self.report(status=409)
|
|
self.assertEqual(report["code"], "busy")
|
|
self.assertEqual(
|
|
report["running"], {"id": restore.id, "title": "Database restore"}
|
|
)
|
|
|
|
async def test_report_lists_config_only_keys(self):
|
|
self.seed_secrets()
|
|
report, _text = await self.report()
|
|
self.assertEqual({k: set(v) for k, v in report["config"].items()}, CONFIG)
|
|
self.assertNotIn("default_path", json.dumps(report["config"]))
|
|
self.assertNotIn("directories", json.dumps(report["config"]))
|
|
self.assertNotIn(
|
|
"civitai_api_key",
|
|
json.dumps(report["config"]).replace("has_civitai_api_key", ""),
|
|
)
|
|
|
|
async def test_report_logging_has_public_directory(self):
|
|
report, _text = await self.report()
|
|
stats = report["logging"]["stats"]
|
|
self.assertFalse(posixpath.isabs(stats["log_directory"]))
|
|
self.assertIn("level", report["logging"]["config"])
|
|
|
|
async def test_report_survives_an_unreadable_viewerjs_provenance(self):
|
|
with mock.patch.object(
|
|
diagnostics_report,
|
|
"VIEWERJS_PROVENANCE",
|
|
os.path.join(self.tmpdir, "nope.json"),
|
|
):
|
|
report, _text = await self.report()
|
|
self.assertEqual(report["viewerjs"], "unknown")
|
|
|
|
async def test_a_build_failure_is_coded_without_details(self):
|
|
with mock.patch.object(
|
|
self.api, "build_report", side_effect=RuntimeError(self.tmpdir)
|
|
):
|
|
with self.assertLogs("prompt_manager", "ERROR"):
|
|
report, text = await self.report(status=500)
|
|
self.assertNotIn(self.tmpdir, text)
|
|
self.assertFalse(report["success"])
|
|
|
|
|
|
class TestLegacyDiagnostics(DiagnosticsReportTestCase):
|
|
async def test_legacy_diagnostics_keeps_its_four_checks(self):
|
|
response = await self.client.request("GET", "/prompt_manager/diagnostics")
|
|
data = await response.json()
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(set(data), {"success", "diagnostics"})
|
|
self.assertEqual(
|
|
set(data["diagnostics"]),
|
|
{"database", "dependencies", "comfyui_output", "image_monitor"},
|
|
)
|
|
report, _text = await self.report()
|
|
self.assertEqual(report["checks"]["database"], data["diagnostics"]["database"])
|
|
self.assertEqual(
|
|
report["checks"]["dependencies"], data["diagnostics"]["dependencies"]
|
|
)
|
|
|
|
async def test_both_routes_build_their_checks_with_one_helper(self):
|
|
with mock.patch.object(
|
|
admin, "legacy_checks", wraps=admin.legacy_checks
|
|
) as spy:
|
|
await self.client.request("GET", "/prompt_manager/diagnostics")
|
|
await self.report()
|
|
self.assertEqual(spy.call_count, 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|