530 lines
21 KiB
Python
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
|
|
import unittest.mock as 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()
|