Files
ComfyAssets-ComfyUI_PromptM…/tests/test_api_library_location.py
T
Vito 8bf759d6f8 feat(system): keep the library in ComfyUI's user folder (#318)
* feat(system): resolve every library path in one place

utils/library_location.py decides where the library lives: the
PROMPT_MANAGER_DATA_DIR folder (validated: absolute, never inside a folder
ComfyUI or Prompt Manager serves), an explicit database.default_path as
before, a node-folder library that was not moved, else ComfyUI's
user/__promptmanager. The database, diagnostics, config.json, logs and the
LoRA image cache all ask it, so a new install never writes into the node
folder and an existing one keeps every path until the user moves it.
"Keep it here" is persisted as library.keep_in_node_folder.

* feat(system): move the library as a verified Activity job

POST /prompt_manager/system/library/move starts "Move library", an
exclusive, non-cancellable job (database/library_move.py). It checks the
target volume's free space, copies the database with SQLite's backup API
from one read snapshot, verifies the copy (integrity check, every workflow
document, the same rows in every table), copies the safety and
pre-migration backups and config.json (owner-only, checked by size and
SHA-256), then publishes under the old file's write lock and switches the
live process (every PromptModel, the review and history observers, the
API's plans and health report). The old prompts.db and config.json are
renamed *.moved-to-<folder> only after the switch. A failure removes
everything it wrote and leaves the old library in use; a marker in the new
folder makes a restart after any crash pick the right library.

GET /prompt_manager/system/library and the Overview's "location" say where
the library lives and why; POST .../keep persists "Keep it here".

* feat(system): show where the library lives and offer the move

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.

* feat(system): close the browser server's log files before its root goes

The logs now live in the library's data folder, which for the browser test
server is <root>/user/__promptmanager; Windows could not remove the temp
root while the server still held prompt_manager.log open. The API test's
node folder is spelled in lower case because public paths are case-folded
on Windows.

* feat(system): show the library location on the legacy admin page

A line under the admin page's statistics names the database path and why it is there (ComfyUI's user folder, PROMPT_MANAGER_DATA_DIR, config.json, or the node folder with a pointer to Move library), read once from GET /prompt_manager/system/library.

* feat(system): prove a fresh host keeps its library in user/__promptmanager

The execution suite starts one pinned CPU ComfyUI with no configuration at all on an OS-chosen port: the library is created in user/__promptmanager (database and logs), nothing but byte code is added to the node folder, the location route reports the default folder, and ComfyUI's /userdata API serves none of it (traversal spellings, a listing, a system-user header).

* feat(system): document where the library lives and how it moves

README: where the library lives, the upgrade path (Move library / Keep it here), PROMPT_MANAGER_DATA_DIR, explicit paths unchanged, and the way back to an older version. TESTING.md: D129 and the library location contract (resolution order, move steps, crash windows, coded refusals, Windows handling).

* wip(system): library move fix round in progress

* fix(system): retry a fallback image link that races the library switch

The image monitor's fallback and recent-prompt links went through
PromptDatabase.link_image_to_prompt, which logs and returns 0 on any
error, so survive_switch never saw a "database is locked" or closed
connection during Move library and the link was lost in both files.
The monitor now links through link_image_to_prompt_or_raise (one
transaction per attempt); link_image_to_prompt keeps its
log-and-return-0 contract for other callers. A link that returns 0 is
no longer logged as "Successfully linked image 0".

A switch that fails after repointing the models now closes the old
file's detached connections instead of leaving them to GC: switch()
adds them to a list the mover owns.

The mocked retry test is replaced by a real-lock race asserting the
image row lands in the new library only.

* fix(system): say which source is the default and where a symlinked library lives

Sources & scans: with no source configured, ComfyUI's output folder now
reads "ComfyUI/output · watching · default" on its own row instead of a
separate note line. GET /settings marks each effective source with
in_comfyui (its public path is relative to the ComfyUI base) so only
those get the "ComfyUI/" prefix; the scan still sends the server's
public path.

Library location: a library in the node folder reads through the node
folder's unresolved place in custom_nodes, so a node folder symlinked to
a checkout outside ComfyUI shows custom_nodes/promptmanager/prompts.db
instead of prompts.db (System library row, Move review, admin Library
line). Never absolute or "..", else the shared _public_path, which
sources and jobs keep using unchanged.
2026-10-24 12:54:09 -07:00

428 lines
18 KiB
Python

"""Library location routes (#249 S6b): where it lives, Move library, Keep it here.
The API's database is a node-folder library inside a fake ComfyUI tree; the
resolver is pointed at that tree (``api._library_resolver``) so nothing here
reads the real environment or the repository's node folder.
"""
import asyncio
import json
import os
import sys
import threading
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 database import library_move as lm # noqa: E402
from database.models import PromptModel # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.config import PromptManagerConfig # noqa: E402
from utils import library_location as ll # 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 = 10
URL = "/prompt_manager/system/library"
class LibraryAPITestCase(AdminAPITestCase):
async def get_application(self):
app = await super().get_application()
# Lower case: a public path is case-folded on Windows (normcase).
self.ext = os.path.join(str(self.comfy_dir), "custom_nodes", "pm")
os.makedirs(self.ext)
self.user = os.path.join(str(self.comfy_dir), "user")
os.makedirs(self.user)
self.target = os.path.join(self.user, "__promptmanager")
sys.modules["folder_paths"].get_user_directory = lambda: self.user
previous = self.api.db
self.api.db = PromptDatabase(os.path.join(self.ext, "prompts.db"))
previous.close_all()
self.api.db.save_prompt("a library prompt")
self.fp = sys.modules["folder_paths"]
self.environ = {}
self.api._library_resolver = lambda: ll.resolve_library(
None, environ=self.environ, extension_root=self.ext, folder_paths=self.fp
)
ll.clear_override()
self.addCleanup(ll.clear_override)
self.addCleanup(PromptModel.close_all_instances)
self._keep = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER
self.addCleanup(
setattr, PromptManagerConfig, "KEEP_LIBRARY_IN_NODE_FOLDER", self._keep
)
PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = False
return app
async def json(self, method, url, status=200, **kwargs):
response = await self.client.request(method, url, **kwargs)
text = await response.text()
self.assertEqual(response.status, status, text)
self.assertNotIn(self.tmpdir, text)
return json.loads(text)
async def finished(self, body):
job = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
return job
class TestLocationRoute(LibraryAPITestCase):
async def test_node_folder_library_offers_the_move(self):
body = await self.json("GET", URL)
self.assertEqual(
body["location"],
{
"path": "custom_nodes/pm/prompts.db",
"source": "node_folder",
"movable": True,
"keep": False,
"offer": True,
"target": "user/__promptmanager",
"problem": None,
"size_bytes": os.path.getsize(self.api.db.model.db_path),
"target_source": "default",
"config_moves": True,
},
)
async def test_overview_carries_the_location(self):
body = await self.json("GET", "/prompt_manager/system/overview")
self.assertEqual(body["location"]["source"], "node_folder")
self.assertTrue(body["location"]["offer"])
async def test_default_folder_library_offers_nothing(self):
self.api._library_resolver = lambda: ll.LibraryLocation(
ll.SOURCE_DEFAULT,
self.api.db.model.db_path,
os.path.join(self.target, "config.json"),
self.target,
self.target,
)
location = (await self.json("GET", URL))["location"]
self.assertEqual((location["source"], location["movable"]), ("default", False))
self.assertFalse(location["offer"])
async def test_a_resolved_file_that_is_not_the_live_database_is_never_moved(self):
self.api._library_resolver = lambda: ll.resolve_library(
None,
environ={},
extension_root=os.path.join(self.tmpdir, "elsewhere"),
folder_paths=self.fp,
)
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["source"], "config")
self.assertFalse(location["movable"])
refused = await self.json("POST", URL + "/move", 400)
self.assertEqual(refused["code"], "not-movable")
class TestLocationDetails(LibraryAPITestCase):
async def test_the_target_and_what_moves_follow_the_environment(self):
data = os.path.join(self.tmpdir, "pm-data")
self.environ = {
ll.DATA_DIR_ENV: data,
ll.CONFIG_PATH_ENV: os.path.join(self.tmpdir, "mine.json"),
}
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["target_source"], "environment")
self.assertFalse(location["config_moves"])
self.assertEqual(location["problem"], "data-dir-empty")
async def test_a_configured_library_inside_the_node_folder_is_flagged(self):
self.api._library_resolver = lambda: ll.LibraryLocation(
ll.SOURCE_CONFIG,
self.api.db.model.db_path,
os.path.join(self.ext, "config.json"),
self.ext,
)
self.api._extension_root = self.ext
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["problem"], "inside-node-folder")
self.assertFalse(location["movable"])
async def test_a_startup_after_a_crash_finishes_the_rename(self):
target_db = os.path.join(self.target, "prompts.db")
os.makedirs(self.target)
with open(target_db, "wb") as handle:
handle.write(b"x")
lm.write_marker(self.target, ll.MARKER_MOVED, self.api.db.model.db_path)
old = self.api.db.model.db_path
self.api.db.model.close_all()
moved = ll.LibraryLocation(
ll.SOURCE_DEFAULT,
target_db,
os.path.join(self.target, "config.json"),
self.target,
self.target,
)
self.api._library_resolver = lambda: moved
with patch.object(self.api.db.model, "db_path", target_db):
renamed = self.api.finish_interrupted_library_move()
self.assertEqual(renamed, ["prompts.db.moved-to-__promptmanager"])
self.assertFalse(os.path.exists(old))
class TestLocationThroughASymlinkedNodeFolder(LibraryAPITestCase):
"""custom_nodes/<name> is a symlink to a checkout outside ComfyUI.
The library location reads relative to ComfyUI through the node folder's
own (unresolved) place in custom_nodes, not as a bare file name; the
shared ``_public_path`` (sources, jobs) keeps resolving symlinks.
"""
async def get_application(self):
app = await super().get_application()
self.real = os.path.join(self.tmpdir, "checkout", "ComfyUI_PromptManager")
os.makedirs(self.real)
self.link = os.path.join(str(self.comfy_dir), "custom_nodes", "promptmanager")
try:
os.symlink(self.real, self.link, target_is_directory=True)
except (OSError, NotImplementedError) as exc:
self.skipTest(f"symlinks not permitted here: {exc}")
previous = self.api.db
self.api.db = PromptDatabase(os.path.join(self.link, "prompts.db"))
previous.close_all()
self.api.db.save_prompt("a library prompt")
self.api._library_resolver = lambda: ll.resolve_library(
None, environ=self.environ, extension_root=self.link, folder_paths=self.fp
)
self.api._extension_root = self.link
return app
async def test_location_names_the_node_folder_in_custom_nodes(self):
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["source"], "node_folder")
self.assertEqual(location["path"], "custom_nodes/promptmanager/prompts.db")
self.assertEqual(location["target"], "user/__promptmanager")
overview = await self.json("GET", "/prompt_manager/system/overview")
self.assertEqual(
overview["location"]["path"], "custom_nodes/promptmanager/prompts.db"
)
async def test_a_resolved_spelling_of_the_database_reads_the_same(self):
from py.api import library_location as routes
resolved = os.path.join(os.path.realpath(self.real), "prompts.db")
self.assertEqual(
routes.node_folder_public_path(resolved, self.link, []),
"custom_nodes/promptmanager/prompts.db",
)
async def test_the_shared_public_path_is_unchanged(self):
self.assertEqual(
self.api._public_path(os.path.join(self.link, "prompts.db")), "prompts.db"
)
class TestNodeFolderPublicPath(LibraryAPITestCase):
def display(self, path, root=None, custom_nodes=()):
from py.api import library_location as routes
return routes.node_folder_public_path(path, root or self.ext, custom_nodes)
async def test_a_plain_node_folder_reads_relative_to_custom_nodes(self):
self.assertEqual(
self.display(os.path.join(self.ext, "prompts.db")),
"custom_nodes/pm/prompts.db",
)
self.assertEqual(
self.display(os.path.join(self.ext, "data", "prompts.db")),
"custom_nodes/pm/data/prompts.db",
)
async def test_a_custom_nodes_entry_from_comfyui_names_the_folder(self):
nodes = os.path.join(self.tmpdir, "more_nodes")
root = os.path.join(nodes, "pm")
self.assertEqual(
self.display(os.path.join(root, "prompts.db"), root, [nodes]),
"more_nodes/pm/prompts.db",
)
async def test_anything_else_is_not_a_node_folder_path(self):
# Outside the node folder, a parent that is no custom_nodes folder, or
# a spelling that climbs out: the caller falls back to _public_path.
self.assertIsNone(self.display(os.path.join(self.user, "prompts.db")))
loose = os.path.join(self.tmpdir, "somewhere", "pm")
self.assertIsNone(self.display(os.path.join(loose, "prompts.db"), loose))
self.assertIsNone(self.display(os.path.join(self.ext, "..", "prompts.db")))
self.assertIsNone(self.display(""))
self.assertIsNone(self.display(None))
async def test_a_library_outside_the_node_folder_keeps_the_public_path(self):
self.api._library_resolver = lambda: ll.LibraryLocation(
ll.SOURCE_DEFAULT,
self.api.db.model.db_path,
os.path.join(self.target, "config.json"),
self.target,
self.target,
)
self.api._extension_root = os.path.join(self.tmpdir, "elsewhere")
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["path"], "custom_nodes/pm/prompts.db")
class TestMoveRoute(LibraryAPITestCase):
async def test_move_runs_as_an_exclusive_job_and_switches_the_live_library(self):
body = await self.json("POST", URL + "/move", 202)
summary = body["job"]
self.assertEqual(summary["kind"], "library_move")
self.assertEqual(summary["title"], "Move library")
self.assertEqual(summary["admission"], "exclusive")
self.assertFalse(summary["cancellable"])
job = await self.finished(body)
self.assertEqual(job.state, "finished", job.status_line)
self.assertEqual(
self.api.db.model.db_path, os.path.join(self.target, "prompts.db")
)
self.assertIn("user/__promptmanager", job.status_line)
self.assertIn("1 prompt", job.status_line)
self.assertNotIn(self.tmpdir, json.dumps(job.report()))
location = (await self.json("GET", URL))["location"]
self.assertEqual(location["source"], "default")
self.assertEqual(location["path"], "user/__promptmanager/prompts.db")
self.assertFalse(location["offer"])
prompts = await self.json("GET", "/prompt_manager/stats")
self.assertEqual(prompts["stats"]["total_prompts"], 1)
self.assertTrue(
os.path.isfile(
os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
)
)
async def test_refusals_are_coded_without_paths(self):
with patch.object(lm, "_free_bytes", return_value=1024):
short = await self.json("POST", URL + "/move", 507)
self.assertEqual(short["code"], "insufficient-space")
os.makedirs(self.target)
with open(os.path.join(self.target, "prompts.db"), "wb") as handle:
handle.write(b"another library")
taken = await self.json("POST", URL + "/move", 409)
self.assertEqual(taken["code"], "target-has-library")
self.assertEqual(self.api._jobs.list(), [])
async def test_move_is_refused_while_another_job_runs_and_blocks_others(self):
gate = asyncio.Event()
async def scanning(job):
await gate.wait()
self.api._jobs.start(
"scan", "Metadata scan", admission="records", runner=scanning
)
busy = await self.json("POST", URL + "/move", 409)
self.assertEqual(busy["code"], "busy")
gate.set()
await asyncio.sleep(0)
for job in self.api._jobs.list():
if job.task is not None:
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
release = threading.Event()
entered = threading.Event()
real = lm.move_library
def gated(*args, **kwargs):
entered.set()
release.wait(WAIT)
return real(*args, **kwargs)
with patch.object(lm, "move_library", side_effect=gated):
body = await self.json("POST", URL + "/move", 202)
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, entered.wait, WAIT)
refused = await self.json("POST", "/prompt_manager/system/health", 409)
self.assertEqual(refused["running"]["title"], "Move library")
release.set()
job = await self.finished(body)
self.assertEqual(job.state, "finished", job.status_line)
async def test_a_failed_move_ends_truthfully_and_keeps_the_old_library(self):
with patch.object(
lm, "move_library", side_effect=lm.MoveFailed("verify-failed")
):
body = await self.json("POST", URL + "/move", 202)
job = await self.finished(body)
self.assertEqual(job.state, "failed")
self.assertIn("did not pass verification", job.status_line)
self.assertIn("still in the node folder", job.status_line)
self.assertEqual(
self.api.db.model.db_path, os.path.join(self.ext, "prompts.db")
)
with patch.object(lm, "move_library", side_effect=OSError("disk gone")):
body = await self.json("POST", URL + "/move", 202)
job = await self.finished(body)
self.assertEqual(job.state, "failed")
self.assertNotIn("disk gone", job.status_line)
self.assertIn("still in the node folder", job.status_line)
async def test_old_file_left_in_place_is_reported(self):
real = lm.retire_old_files
with patch.object(lm, "retire_old_files", return_value=([], "old-in-use")):
body = await self.json("POST", URL + "/move", 202)
job = await self.finished(body)
self.assertEqual(job.state, "finished")
self.assertIn("left in the node folder", job.status_line)
self.assertTrue(callable(real))
class TestKeepRoute(LibraryAPITestCase):
async def test_keep_it_here_is_persisted_and_stops_the_offer(self):
body = await self.json("POST", URL + "/keep", 200, json={"keep": True})
self.assertTrue(body["location"]["keep"])
self.assertFalse(body["location"]["offer"])
self.assertTrue(body["location"]["movable"])
with open(self.config_path, encoding="utf-8") as handle:
saved = json.load(handle)
self.assertEqual(saved["library"], {"keep_in_node_folder": True})
self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
undone = await self.json("POST", URL + "/keep", 200, json={"keep": False})
self.assertTrue(undone["location"]["offer"])
async def test_keep_needs_a_boolean(self):
for payload in ({"keep": "yes"}, {}, [True]):
bad = await self.json("POST", URL + "/keep", 400, json=payload)
self.assertEqual(bad["code"], "invalid-request")
self.assertFalse(os.path.exists(self.config_path))
async def test_unwritable_config_reports_not_saved(self):
with patch("py.config_store._write_atomically", side_effect=OSError("ro")):
body = await self.json("POST", URL + "/keep", 500, json={"keep": True})
self.assertEqual(body["code"], "config-write-failed")
self.assertFalse(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER)
class TestLegacyRoutesDuringAMove(LibraryAPITestCase):
async def test_legacy_restore_is_refused_while_the_move_runs(self):
from aiohttp import FormData
release = threading.Event()
entered = threading.Event()
real = lm.move_library
def gated(*args, **kwargs):
entered.set()
release.wait(WAIT)
return real(*args, **kwargs)
with patch.object(lm, "move_library", side_effect=gated):
body = await self.json("POST", URL + "/move", 202)
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, entered.wait, WAIT)
form = FormData()
form.add_field("database_file", b"SQLite format 3\x00", filename="x.db")
response = await self.client.request(
"POST", "/prompt_manager/restore", data=form
)
self.assertEqual(response.status, 409, await response.text())
release.set()
await self.finished(body)