Files
ComfyAssets-ComfyUI_PromptM…/tests/test_library_move.py
T
Vito 4040007264 fix(workspace): keep open pages valid across a restart and recover a changed source (#430) (#434)
The gallery root scope and the Images listing cursor were keyed with a
per-process random key, so every ComfyUI restart made an open tab's
scopes answer 409 source-changed (blank Tag image preview, a Generate
that did nothing) and its cursors fail Load more.

- utils/workspace_keys.py: one random secret in workspace-scope.key
  beside the database (owner-only, never config.json or the database),
  two derived keys; per-process fallback with a warning. Move library
  copies it.
- Cursor v2 binds the listing scope in its MAC instead of sending an
  unkeyed hash of the root paths; an unverifiable cursor answers 400
  cursor-stale and Images offers Refresh instead of a dead Retry.
- POST /prompt_manager/images/rescope maps a stale target to its
  current root only when the held scope was issued for it.
- Images finds a renumbered file again from the inspector, the viewer
  and the Tag image dialog (the dialog stays open, draft kept); nothing
  else is kept by an old root_index:path key.
- Tag image dialog: a state that blocks Save shows its reason above the
  actions with Refresh, disabled actions look disabled, and an empty
  preview says "Preview unavailable" (D147).
2026-10-28 23:06:36 +00:00

530 lines
21 KiB
Python

"""Move library (#249 S6b): copy, verify, switch, then retire the old files.
Each test builds a node-folder library (a real PromptDatabase with prompts,
images and workflow documents, a config.json holding a key, a pooled safety
backup and the pre-migration backup) and a ComfyUI user folder in a temp
directory, then moves it with ``database.library_move``. Failures and crash
windows are injected through the move's ``checkpoint`` hook: an ``Exception``
is a failure the move cleans up after, ``SimulatedCrash`` (a BaseException)
stops it dead like a killed process.
"""
import hashlib
import json
import os
import sqlite3
import sys
import tempfile
import threading
import types
import unittest
from unittest import mock
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.execution_history import ( # noqa: E402
notify_history_restore,
subscribe_history_restore,
)
from database.models import PromptModel # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from database.review_changes import ( # noqa: E402
notify_review_change,
subscribe_review_changes,
)
from tests.open_handles import assert_closed, open_handles_under # noqa: E402
from utils import library_location as ll # noqa: E402
KEY = "civitai-secret-key"
POOL = "prompts.db.backup_20261001_101010_000001"
PRE = "prompts.db.pre-workflow-storage_20261002_101010_000002"
class SimulatedCrash(BaseException):
"""The process died here: nothing after this point runs."""
def _digest(path):
with open(path, "rb") as handle:
return hashlib.sha256(handle.read()).hexdigest()
def _rows(path):
"""Every table's rows of ``path`` (read-only), to compare libraries."""
uri = "file:" + path.replace(os.sep, "/") + "?mode=ro"
conn = sqlite3.connect(uri, uri=True)
try:
tables = [
row[0]
for row in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' "
"AND name NOT LIKE 'sqlite_%' ORDER BY name"
)
]
return {
table: sorted(
repr(tuple(row)) for row in conn.execute(f'SELECT * FROM "{table}"')
)
for table in tables
}
finally:
conn.close()
def _workflow(seed):
return {
"nodes": [
{"id": n, "type": f"Node{seed}", "widgets": [seed, n]} for n in range(30)
]
}
class _MoveCase(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.base = os.path.realpath(self._tmp.name)
self.comfy = os.path.join(self.base, "ComfyUI")
self.ext = os.path.join(self.comfy, "custom_nodes", "ComfyUI_PromptManager")
os.makedirs(self.ext)
self.fp = types.SimpleNamespace()
for name in ("output", "input", "temp", "user"):
folder = os.path.join(self.comfy, name)
os.makedirs(folder)
setattr(self.fp, f"get_{name}_directory", lambda f=folder: f)
self.target = os.path.join(self.comfy, "user", "__promptmanager")
self.old_db = os.path.join(self.ext, "prompts.db")
self.new_db = os.path.join(self.target, "prompts.db")
ll.clear_override()
self.db = PromptDatabase(self.old_db)
for index in range(3):
prompt_id = self.db.save_prompt(f"a cat number {index}", category="cats")
self.db.link_image_to_prompt(
str(prompt_id),
os.path.join(self.comfy, "output", f"cat{index}.png"),
{"workflow": _workflow(index % 2), "prompt": {"1": {"seed": index}}},
)
self.db.model.close_all()
self.config = os.path.join(self.ext, "config.json")
with open(self.config, "w", encoding="utf-8") as handle:
json.dump(
{"integrations": {"lora_manager": {"civitai_api_key": KEY}}}, handle
)
for name, stamp in ((POOL, 1_700_000_000), (PRE, 1_700_000_100)):
path = os.path.join(self.ext, name)
source = sqlite3.connect(self.old_db)
copy = sqlite3.connect(path)
try:
source.backup(copy)
finally:
copy.close()
source.close()
os.utime(path, (stamp, stamp))
self.before = {
"rows": _rows(self.old_db),
"config": _digest(self.config),
"pool": _digest(os.path.join(self.ext, POOL)),
"pre": _digest(os.path.join(self.ext, PRE)),
}
self.config_lock = threading.RLock()
def tearDown(self):
import gc
gc.collect()
PromptModel.close_all_instances()
ll.clear_override()
try:
assert_closed(self, self.ext, self.target)
finally:
self._tmp.cleanup()
def location(self, **kwargs):
return ll.resolve_library(
None, environ={}, extension_root=self.ext, folder_paths=self.fp, **kwargs
)
def move(self, **kwargs):
plan = lm.plan_move(self.location())
kwargs.setdefault("config_lock", self.config_lock)
return lm.move_library(plan, **kwargs)
def assert_old_library_untouched(self):
self.assertTrue(os.path.isfile(self.old_db))
self.assertEqual(_rows(self.old_db), self.before["rows"])
self.assertEqual(_digest(self.config), self.before["config"])
self.assertEqual(_digest(os.path.join(self.ext, POOL)), self.before["pool"])
self.assertEqual(_digest(os.path.join(self.ext, PRE)), self.before["pre"])
self.assertEqual(
os.path.normcase(self.db.model.db_path), os.path.normcase(self.old_db)
)
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
def assert_target_clean(self):
leftovers = os.listdir(self.target) if os.path.isdir(self.target) else []
self.assertEqual(leftovers, [])
class TestMove(_MoveCase):
def test_move_copies_verifies_switches_and_retires_the_old_files(self):
result = self.move()
self.assertEqual(_rows(self.new_db), self.before["rows"])
self.assertEqual(result.counts["prompts"], 3)
self.assertEqual(result.counts["generated_images"], 3)
self.assertEqual(
result.counts["workflow_documents"], 5
) # 2 workflows, 3 graphs
self.assertEqual(
sorted(result.files), sorted(["config.json", POOL, PRE, "prompts.db"])
)
self.assertIsNone(result.retire_problem)
# The live database object now reads and writes the new file.
self.assertEqual(self.db.model.db_path, self.new_db)
self.db.save_prompt("written after the move")
self.assertEqual(len(_rows(self.new_db)["prompts"]), 4)
image = self.db.get_prompt_images(
str(self.db.search_prompts(text="number 1")[0]["id"])
)[0]
self.assertEqual(image["workflow_data"], _workflow(1))
# config.json, the pooled backup and the pre-migration backup: same bytes.
self.assertEqual(
_digest(os.path.join(self.target, "config.json")), self.before["config"]
)
self.assertEqual(_digest(os.path.join(self.target, POOL)), self.before["pool"])
self.assertEqual(_digest(os.path.join(self.target, PRE)), self.before["pre"])
self.assertEqual(
int(os.path.getmtime(os.path.join(self.target, POOL))), 1_700_000_000
)
# The old files are renamed, never deleted; the backups stay.
renamed = sorted(os.listdir(self.ext))
self.assertIn("prompts.db.moved-to-__promptmanager", renamed)
self.assertIn("config.json.moved-to-__promptmanager", renamed)
self.assertNotIn("prompts.db", renamed)
self.assertNotIn("config.json", renamed)
self.assertFalse(any(name.endswith(("-wal", "-shm")) for name in renamed))
self.assertEqual(
_rows(os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")),
self.before["rows"],
)
self.assertEqual(ll.read_marker(self.target)["state"], ll.MARKER_MOVED)
def test_the_workspace_key_moves_with_the_library(self):
# #430: open pages keep their scopes after the restart that follows.
from utils.workspace_keys import KEY_FILE_NAME, load_install_secret
secret = load_install_secret(self.ext)
result = self.move()
self.assertIn(KEY_FILE_NAME, result.files)
self.assertEqual(load_install_secret(self.target), secret)
if os.name == "posix":
mode = os.stat(os.path.join(self.target, KEY_FILE_NAME)).st_mode
self.assertEqual(mode & 0o777, 0o600)
def test_a_key_that_cannot_be_copied_never_fails_the_move(self):
from utils.workspace_keys import KEY_FILE_NAME, load_install_secret
load_install_secret(self.ext)
real = lm._hash_copy
def failing(source, destination):
if os.path.basename(source) == KEY_FILE_NAME:
raise OSError("disk error")
return real(source, destination)
with mock.patch.object(lm, "_hash_copy", failing):
with self.assertLogs("prompt_manager.database.library_move", "WARNING"):
result = self.move()
self.assertNotIn(KEY_FILE_NAME, result.files)
self.assertEqual(_rows(self.new_db), self.before["rows"])
self.assertFalse(
any(n.startswith(lm.TEMP_PREFIX) for n in os.listdir(self.target))
)
def test_new_files_are_private(self):
if os.name != "posix":
self.skipTest("POSIX permissions")
self.move()
self.assertEqual(os.stat(self.target).st_mode & 0o777, 0o700)
for name in ("prompts.db", "config.json", POOL, PRE, ll.MARKER_NAME):
mode = os.stat(os.path.join(self.target, name)).st_mode & 0o777
self.assertEqual(mode, 0o600, name)
def test_restart_uses_the_new_library_and_never_the_renamed_file(self):
self.move()
self.assertEqual(self.location().db_path, self.new_db)
ll.clear_override()
restarted = self.location()
self.assertEqual(restarted.source, ll.SOURCE_DEFAULT)
self.assertEqual(restarted.db_path, self.new_db)
self.assertEqual(
restarted.config_path, os.path.join(self.target, "config.json")
)
def test_every_open_database_object_is_repointed(self):
node = PromptDatabase(self.old_db)
monitor = PromptDatabase(self.old_db)
monitor.get_prompt_by_id(1)
other = PromptDatabase(os.path.join(self.base, "other.db"))
self.move()
self.assertEqual(node.model.db_path, self.new_db)
self.assertEqual(monitor.model.db_path, self.new_db)
self.assertEqual(monitor.get_prompt_by_id(1)["text"], "a cat number 0")
self.assertEqual(other.model.db_path, os.path.join(self.base, "other.db"))
assert_closed(
self, os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
)
def test_change_observers_follow_the_library(self):
reviews, restores = [], []
unsubscribe_review = subscribe_review_changes(
self.old_db, lambda: reviews.append(1)
)
unsubscribe_restore = subscribe_history_restore(
self.old_db, lambda: restores.append(1)
)
self.addCleanup(unsubscribe_review)
self.addCleanup(unsubscribe_restore)
self.move()
self.assertEqual(restores, [1]) # the history readers start a new epoch
notify_review_change(self.new_db)
notify_history_restore(self.new_db)
self.assertEqual(reviews, [1])
self.assertEqual(restores, [1, 1])
def test_files_are_only_renamed_inside_one_folder(self):
real_replace, real_rename = os.replace, os.rename
moves = []
def replace(src, dst, *args, **kwargs):
moves.append((src, dst))
return real_replace(src, dst, *args, **kwargs)
def rename(src, dst, *args, **kwargs):
moves.append((src, dst))
return real_rename(src, dst, *args, **kwargs)
with (
mock.patch.object(lm.os, "replace", replace),
mock.patch.object(lm.os, "rename", rename),
):
self.move()
self.assertTrue(moves)
for src, dst in moves:
self.assertEqual(os.path.dirname(src), os.path.dirname(dst), (src, dst))
def test_a_write_during_the_copy_is_copied_again(self):
writes = []
def checkpoint(step):
if step == "files" and not writes:
writes.append(PromptDatabase(self.old_db).save_prompt("late write"))
progress = []
self.move(checkpoint=checkpoint, progress=progress.append)
self.assertIn("late write", repr(_rows(self.new_db)["prompts"]))
self.assertTrue(any("changed" in line for line in progress), progress)
def test_a_library_that_keeps_changing_is_not_moved(self):
def checkpoint(step):
if step == "files":
PromptDatabase(self.old_db).save_prompt(f"write {os.urandom(4).hex()}")
with self.assertRaises(lm.MoveFailed) as caught:
self.move(checkpoint=checkpoint, attempts=2)
self.assertEqual(caught.exception.code, "library-busy")
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
self.assert_target_clean()
def test_old_file_still_open_elsewhere_is_left_in_place(self):
reader = sqlite3.connect(self.old_db)
self.addCleanup(reader.close)
reader.execute("BEGIN")
reader.execute("SELECT COUNT(*) FROM prompts").fetchone()
PromptDatabase(self.old_db).save_prompt("in the wal")
result = self.move()
self.assertEqual(result.retire_problem, "old-in-use")
self.assertTrue(os.path.isfile(self.old_db))
self.assertEqual(self.db.model.db_path, self.new_db)
reader.rollback()
reader.close()
ll.clear_override()
self.assertEqual(self.location().source, ll.SOURCE_DEFAULT)
class TestFailures(_MoveCase):
def _fail_at(self, step):
def checkpoint(name):
if name == step:
raise RuntimeError(f"failure at {step}")
return checkpoint
def test_a_failure_at_any_step_leaves_the_old_library_in_use_and_untouched(self):
for step in ("copied", "verified", "files", "locked", "pending", "published"):
with self.subTest(step=step):
with self.assertRaises(RuntimeError):
self.move(checkpoint=self._fail_at(step))
self.assert_old_library_untouched()
self.assert_target_clean()
assert_closed(self, self.target)
def test_a_copy_that_does_not_verify_is_refused(self):
def checkpoint(step):
if step == "copied":
(temp,) = [
name for name in os.listdir(self.target) if name.endswith(".db")
]
with open(os.path.join(self.target, temp), "r+b") as handle:
handle.truncate(4096)
with self.assertRaises(lm.MoveFailed) as caught:
self.move(checkpoint=checkpoint)
self.assertEqual(caught.exception.code, "verify-failed")
self.assertNotIn(self.base, caught.exception.message)
self.assert_old_library_untouched()
self.assert_target_clean()
def test_a_copied_file_that_differs_is_refused(self):
def checkpoint(step):
if step == "files":
for name in os.listdir(self.target):
if name.endswith(".bak"):
with open(os.path.join(self.target, name), "ab") as handle:
handle.write(b"x")
with self.assertRaises(lm.MoveFailed) as caught:
self.move(checkpoint=checkpoint)
self.assertEqual(caught.exception.code, "copy-failed")
self.assert_old_library_untouched()
self.assert_target_clean()
def test_not_movable_locations_are_refused(self):
os.makedirs(self.target)
with open(self.new_db, "wb") as handle:
handle.write(b"someone else's library")
with self.assertRaises(lm.MoveFailed) as caught:
lm.plan_move(self.location())
self.assertEqual(caught.exception.code, "target-has-library")
os.remove(self.new_db)
default = ll.LibraryLocation(
ll.SOURCE_DEFAULT, self.new_db, "c", self.target, self.target
)
with self.assertRaises(lm.MoveFailed) as caught:
lm.plan_move(default)
self.assertEqual(caught.exception.code, "not-movable")
def test_target_that_is_the_node_folder_is_refused(self):
link = os.path.join(self.base, "link-to-node-folder")
try:
os.symlink(self.ext, link)
except (OSError, NotImplementedError):
self.skipTest("symlinks are not available")
location = ll.LibraryLocation(
ll.SOURCE_NODE_FOLDER, self.old_db, self.config, self.ext, link
)
with self.assertRaises(lm.MoveFailed) as caught:
lm.plan_move(location)
self.assertEqual(caught.exception.code, "same-folder")
def test_disk_check_measures_the_target_volume(self):
measured = []
def free(path):
measured.append(path)
return 1024
with mock.patch.object(lm, "_free_bytes", free):
plan = lm.plan_move(self.location())
with self.assertRaises(lm.MoveFailed) as caught:
lm.move_library(plan, config_lock=self.config_lock)
self.assertEqual(caught.exception.code, "insufficient-space")
self.assertTrue(measured)
for path in measured:
self.assertTrue(
os.path.normcase(self.comfy) in os.path.normcase(path)
and "custom_nodes" not in path,
path,
)
self.assertGreater(plan.required_bytes, os.path.getsize(self.old_db))
self.assert_old_library_untouched()
class TestCrashWindows(_MoveCase):
"""A crash between steps never loses the library or picks a stale one."""
def _crash_at(self, step):
def checkpoint(name):
if name == step:
raise SimulatedCrash(step)
return checkpoint
def _restart(self):
PromptModel.close_all_instances()
ll.clear_override()
return self.location()
def test_crash_before_the_switch_keeps_the_old_library_and_a_retry_succeeds(self):
for step in ("copied", "files", "pending", "published"):
with self.subTest(step=step):
with self.assertRaises(SimulatedCrash):
self.move(checkpoint=self._crash_at(step))
restarted = self._restart()
self.assertEqual(restarted.source, ll.SOURCE_NODE_FOLDER)
self.assertEqual(restarted.db_path, self.old_db)
self.assertEqual(_rows(self.old_db), self.before["rows"])
result = self.move()
self.assertEqual(_rows(self.new_db), self.before["rows"])
self.assertIsNone(result.retire_problem)
self.assertFalse(
[name for name in os.listdir(self.target) if name.startswith(".pm-move-")]
)
def test_crash_after_the_marker_uses_the_new_library(self):
for step in ("marked", "switched"):
with self.subTest(step=step):
with self.assertRaises(SimulatedCrash):
self.move(checkpoint=self._crash_at(step))
restarted = self._restart()
self.assertEqual(restarted.source, ll.SOURCE_DEFAULT)
self.assertEqual(restarted.db_path, self.new_db)
self.assertEqual(_rows(self.new_db), self.before["rows"])
# The old file was not renamed yet and is never picked.
self.assertTrue(os.path.isfile(self.old_db))
# Reset for the next window: back to the node folder only.
for name in os.listdir(self.target):
os.remove(os.path.join(self.target, name))
def test_crash_after_the_rename(self):
with self.assertRaises(SimulatedCrash):
self.move(checkpoint=self._crash_at("retired"))
restarted = self._restart()
self.assertEqual(restarted.db_path, self.new_db)
self.assertFalse(os.path.exists(self.old_db))
class TestEnvironmentTarget(_MoveCase):
def test_move_into_the_environment_folder(self):
data = os.path.join(self.base, "pm-data")
environ = {ll.DATA_DIR_ENV: data}
location = ll.resolve_library(
None, environ=environ, extension_root=self.ext, folder_paths=self.fp
)
self.assertEqual(location.problem, "data-dir-empty")
result = lm.move_library(lm.plan_move(location), config_lock=self.config_lock)
self.assertEqual(result.location.source, ll.SOURCE_ENVIRONMENT)
self.assertEqual(_rows(os.path.join(data, "prompts.db")), self.before["rows"])
ll.clear_override()
self.assertEqual(
ll.resolve_library(
None, environ=environ, extension_root=self.ext, folder_paths=self.fp
).source,
ll.SOURCE_ENVIRONMENT,
)
self.assertEqual(open_handles_under(self.old_db), [])
if __name__ == "__main__":
unittest.main()