fix(maintenance): never remove the kept copy; move records before unlink
Final review of #313: - C1: the duplicate find tells files apart by (st_dev, st_ino), so a source listed twice, a symlinked source inside another and a hardlink are one file. The preview and the job refuse a removal that os.path.samestat matches with its kept copy. - I1: each file's records move to its verified kept copy in one short transaction before that file is unlinked; a failed move deletes nothing and a failed unlink leaves the records on the identical kept copy. The status line counts files removed and not removed. - I2: links are looked up under every spelling of the source (the real path, the configured and monitored folders, ComfyUI's output). - I5: image tags join the kept copy as their union and the removed copy's pending, undecided Needs review items are settled; the plan says so. - I3: the VACUUM disk check measures SQLite's real temporary folder. - M1: the token is looked at again with no await before the job starts. - M2: the backup decision is made once and handed to the job. - M3: backups are linked into place; startup sweeps interrupted temps. - M4: Optimize runs alone (new "exclusive" admission) and its plan says the database is locked while VACUUM runs. - M5: the reviewed duplicate merge is strict: collisions and Library titles are protected, and errors roll back. The legacy cleanup is unchanged. - M7: 2,000 keep entries within a 1 MiB body. M8: event waits in tests. - M9: the preview binds files by size, mtime and inode without hashing; the job hashes right before each unlink.
This commit is contained in:
Vito Sansevero
@@ -126,6 +126,14 @@ class DuplicateFilesTestCase(ApplyTestCase):
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
return plan, cleanup
|
||||
|
||||
async def until(self, predicate):
|
||||
"""Wait for a server-side state (a frame at a time), never a fixed timer."""
|
||||
deadline = asyncio.get_running_loop().time() + WAIT
|
||||
while not predicate():
|
||||
if asyncio.get_running_loop().time() > deadline:
|
||||
self.fail("the expected server state never appeared")
|
||||
await asyncio.sleep(0)
|
||||
|
||||
def gate_hashing(self):
|
||||
gate = threading.Event()
|
||||
calls = []
|
||||
@@ -159,7 +167,8 @@ class TestFindJob(DuplicateFilesTestCase):
|
||||
self.assertEqual(summary["counters"], {"groups_found": 0})
|
||||
# The legacy gallery attaches to the same job, and its status says so.
|
||||
legacy = asyncio.ensure_future(self.client.request("POST", STREAM))
|
||||
await asyncio.sleep(0.1)
|
||||
job = self.api._jobs.get(summary["id"])
|
||||
await self.until(lambda: job.subscriber_count > 0)
|
||||
status = await (await self.client.request("GET", STATUS)).json()
|
||||
self.assertTrue(status["running"])
|
||||
gate.set()
|
||||
@@ -220,7 +229,8 @@ class TestFindJob(DuplicateFilesTestCase):
|
||||
response = await self.client.request("POST", FIND)
|
||||
job_id = (await response.json())["job"]["id"]
|
||||
stream = asyncio.ensure_future(self.client.request("POST", STREAM))
|
||||
await asyncio.sleep(0.1)
|
||||
found = self.api._jobs.get(job_id)
|
||||
await self.until(lambda: found.subscriber_count > 0)
|
||||
cancel = await self.client.request(
|
||||
"POST", f"/prompt_manager/jobs/{job_id}/cancel"
|
||||
)
|
||||
@@ -512,7 +522,11 @@ class TestCleanupJob(DuplicateFilesTestCase):
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
self.assertTrue(cleanup.status_line.startswith("Removed 0 duplicate files · "))
|
||||
self.assertTrue(
|
||||
cleanup.status_line.startswith(
|
||||
"Removed 0 duplicate files, 2 not removed · "
|
||||
)
|
||||
)
|
||||
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
|
||||
self.assertEqual(rows["b.png"], duplicate_files.KEEPER_CHANGED)
|
||||
|
||||
@@ -544,8 +558,8 @@ class TestCleanupJob(DuplicateFilesTestCase):
|
||||
)
|
||||
self.assertEqual(
|
||||
cleanup.status_line,
|
||||
"Stopped after an error · Removed 1 duplicate file · Safety backup "
|
||||
+ self.backups()[0],
|
||||
"Stopped after an error · Removed 1 duplicate file, 1 not removed · "
|
||||
"Safety backup " + self.backups()[0],
|
||||
)
|
||||
report = json.dumps(cleanup.report())
|
||||
self.assertNotIn("Permission denied", report)
|
||||
@@ -660,7 +674,7 @@ class TestCleanupJob(DuplicateFilesTestCase):
|
||||
"Moved 2 database links to the kept copies", "\n".join(cleanup.log)
|
||||
)
|
||||
|
||||
async def test_a_failed_removal_keeps_its_links(self):
|
||||
async def test_a_failed_removal_leaves_its_links_on_the_kept_copy(self):
|
||||
self.trio()
|
||||
prompt = self.save("Stays on c")
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "c.png"))
|
||||
@@ -677,8 +691,10 @@ class TestCleanupJob(DuplicateFilesTestCase):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
# Its records moved to the identical kept copy before the unlink failed.
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertTrue(image["image_path"].endswith("c.png"))
|
||||
self.assertTrue(image["image_path"].endswith("a.png"))
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
|
||||
async def test_cleanup_and_find_exclude_each_other(self):
|
||||
self.trio()
|
||||
@@ -772,7 +788,7 @@ class TestLegacyDuplicateRoutes(DuplicateFilesTestCase):
|
||||
legacy = asyncio.ensure_future(
|
||||
self.client.request("GET", "/prompt_manager/scan_duplicates")
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
await self.until(lambda: self.api._jobs.active("duplicates") is not None)
|
||||
[hold] = self.api._jobs.list()
|
||||
self.assertEqual(
|
||||
(hold.kind, hold.title, hold.admission, hold.cancellable),
|
||||
@@ -942,14 +958,14 @@ class TestReviewHelpers(DuplicateFilesTestCase):
|
||||
self.assertEqual(self.backups(), [])
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
|
||||
async def test_a_link_move_failure_is_reported(self):
|
||||
async def test_a_records_move_failure_deletes_nothing(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
with (
|
||||
patch.object(
|
||||
self.api.db,
|
||||
"move_file_links",
|
||||
"move_copy_records",
|
||||
side_effect=sqlite3.OperationalError("x"),
|
||||
),
|
||||
self.assertLogs("prompt_manager", "ERROR"),
|
||||
@@ -959,6 +975,8 @@ class TestReviewHelpers(DuplicateFilesTestCase):
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
self.assertEqual(cleanup.state, "failed")
|
||||
self.assertEqual(
|
||||
cleanup.outcomes["errors"][-1]["detail"], duplicate_files.LINKS_FAILED
|
||||
cleanup.outcomes["errors"][-1]["detail"], duplicate_files.RECORDS_FAILED
|
||||
)
|
||||
self.assertEqual(cleanup.counters["removed"], 2)
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
"""Duplicate-file removal safety (#249 System S6 final review: C1, I1, I2, I5, M9).
|
||||
|
||||
- C1: one physical file reached through two listed sources, a symlinked
|
||||
source or a hardlink is never offered or removed as a copy of itself.
|
||||
- I1: each file's records move to its verified kept copy before that file
|
||||
is unlinked, so no database row ever names a deleted file.
|
||||
- I2: links stored under another spelling of the source move too.
|
||||
- I5: image tags move to the kept copy (their union) and the removed copy's
|
||||
pending Needs review items are settled.
|
||||
- M9: the preview binds files by size, mtime and inode without hashing.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from tests import test_api_duplicate_files as dup_tests
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
import test_api_duplicate_files as dup_tests
|
||||
|
||||
from py.api import duplicate_files # noqa: E402
|
||||
from py.config import GalleryConfig # noqa: E402
|
||||
|
||||
WAIT = 10
|
||||
|
||||
|
||||
def _locator(path):
|
||||
return os.path.normcase(os.path.realpath(str(path)))
|
||||
|
||||
|
||||
class SafetyTestCase(dup_tests.DuplicateFilesTestCase):
|
||||
def need_symlinks(self):
|
||||
if not dup_tests.TestReview._can_symlink(self.tmpdir): # pragma: no cover
|
||||
self.skipTest("symlinks are not available")
|
||||
|
||||
def linked_files(self):
|
||||
"""Every image record's path that names a file that does not exist."""
|
||||
with self.api.db.model.get_connection() as conn:
|
||||
rows = conn.execute("SELECT image_path FROM generated_images").fetchall()
|
||||
return [row[0] for row in rows if not os.path.exists(row[0])]
|
||||
|
||||
async def apply_plan(self, plan):
|
||||
body = await self.apply(plan["token"])
|
||||
cleanup = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(cleanup.task), WAIT)
|
||||
return cleanup
|
||||
|
||||
|
||||
class TestOnePhysicalFile(SafetyTestCase):
|
||||
"""C1: the kept copy survives every alias of itself."""
|
||||
|
||||
async def test_a_folder_listed_twice_never_loses_its_only_copy(self):
|
||||
self.need_symlinks()
|
||||
only = self.write("a.png", b"the only copy")
|
||||
alias = Path(self.tmpdir) / "Output-alias"
|
||||
os.symlink(str(self.output_dir), str(alias))
|
||||
with patch.object(
|
||||
self.api, "_get_all_output_dirs", return_value=[self.output_dir, alias]
|
||||
):
|
||||
job = await self.find()
|
||||
page = await self.groups(job.id)
|
||||
self.assertEqual(page["total"], 0)
|
||||
self.assertEqual(only.read_bytes(), b"the only copy")
|
||||
|
||||
async def test_a_symlinked_source_inside_another_is_one_file(self):
|
||||
self.need_symlinks()
|
||||
only = self.write("sub/c.png", b"one file")
|
||||
linked = self.output_dir / "linked" # walked by neither as a link
|
||||
os.symlink(str(self.output_dir / "sub"), str(linked))
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [
|
||||
self._root(self.output_dir),
|
||||
str(linked),
|
||||
]
|
||||
self.assertEqual(len(self.api._get_all_output_dirs()), 2)
|
||||
job = await self.find()
|
||||
self.assertEqual((await self.groups(job.id))["total"], 0)
|
||||
self.assertTrue(only.is_file())
|
||||
|
||||
async def test_a_hardlink_is_never_a_copy_to_remove(self):
|
||||
first = self.write("a.png", b"same bytes", mtime=1)
|
||||
os.link(str(first), str(self.output_dir / "b.png"))
|
||||
self.write("c.png", b"same bytes", mtime=3)
|
||||
job, group = await self.found_group()
|
||||
self.assertEqual(len(group["files"]), 2) # a.png and b.png are one file
|
||||
# c.png becomes a hardlink of the keeper after the find: protected.
|
||||
(self.output_dir / "c.png").unlink()
|
||||
os.link(str(first), str(self.output_dir / "c.png"))
|
||||
kept = group["files"][0]["path"]
|
||||
plan = await self.review(job.id, [self.keep_of(group, kept)])
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"], ["1 file is the same file as its kept copy"]
|
||||
)
|
||||
|
||||
async def test_the_job_refuses_a_file_that_became_its_kept_copy(self):
|
||||
first = self.write("a.png", b"same bytes", mtime=1)
|
||||
self.write("b.png", b"same bytes", mtime=2)
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
(self.output_dir / "b.png").unlink()
|
||||
os.link(str(first), str(self.output_dir / "b.png"))
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertEqual(cleanup.counters["removed"], 0)
|
||||
rows = {row["path"]: row["detail"] for row in cleanup.outcomes["kept"]}
|
||||
self.assertEqual(rows["b.png"], duplicate_files.SAME_FILE)
|
||||
self.assertTrue((self.output_dir / "b.png").is_file())
|
||||
|
||||
|
||||
class TestRecordsMoveBeforeTheUnlink(SafetyTestCase):
|
||||
"""I1: a deleted file is never named by a database row."""
|
||||
|
||||
async def test_a_failure_between_two_files_leaves_no_row_on_a_deleted_file(self):
|
||||
self.write("a.png", b"same", mtime=1)
|
||||
self.write("b.png", b"same", mtime=2)
|
||||
self.write("c.png", b"same", mtime=3)
|
||||
prompt = self.save("Linked to every copy")
|
||||
for name in ("b.png", "c.png"):
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / name))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = self.api.db.move_copy_records
|
||||
calls = []
|
||||
|
||||
def move(*args, **kwargs):
|
||||
calls.append(args)
|
||||
if len(calls) == 2:
|
||||
raise sqlite3.OperationalError("database is locked")
|
||||
return real(*args, **kwargs)
|
||||
|
||||
with patch.object(self.api.db, "move_copy_records", side_effect=move):
|
||||
with self.assertLogs("prompt_manager", "ERROR"):
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
self.assertEqual(cleanup.counters["removed"], 1)
|
||||
self.assertEqual(cleanup.state, "failed")
|
||||
[error] = cleanup.outcomes["errors"]
|
||||
self.assertEqual(error["detail"], duplicate_files.RECORDS_FAILED)
|
||||
self.assertNotIn("Missing-file", cleanup.status_line)
|
||||
self.assertEqual(
|
||||
cleanup.status_line.split(" · ")[:2],
|
||||
["Stopped after an error", "Removed 1 duplicate file, 1 not removed"],
|
||||
)
|
||||
|
||||
async def test_a_failed_unlink_leaves_its_links_on_the_kept_copy(self):
|
||||
self.trio()
|
||||
prompt = self.save("On c")
|
||||
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "sub" / "c.png"))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
real = duplicate_files._unlink
|
||||
|
||||
def unlink(path):
|
||||
if os.path.basename(path) == "c.png":
|
||||
raise PermissionError(13, "denied")
|
||||
real(path)
|
||||
|
||||
with patch.object(duplicate_files, "_unlink", side_effect=unlink):
|
||||
with self.assertLogs("prompt_manager", "WARNING"):
|
||||
cleanup = await self.apply_plan(plan)
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertEqual(os.path.basename(image["image_path"]), "a.png")
|
||||
self.assertTrue((self.output_dir / "sub" / "c.png").is_file())
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
self.assertEqual(
|
||||
cleanup.outcomes["errors"][0]["detail"], duplicate_files.LOCKED_FILE
|
||||
)
|
||||
|
||||
|
||||
class TestEverySpelling(SafetyTestCase):
|
||||
"""I2: links under the monitor's spelling of a source move too."""
|
||||
|
||||
async def test_links_under_a_symlinked_source_spelling_move(self):
|
||||
self.need_symlinks()
|
||||
self.trio()
|
||||
spelled = Path(self.comfy_dir) / "out-link"
|
||||
os.symlink(str(self.output_dir), str(spelled))
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(spelled)]
|
||||
prompt = self.save("Linked through the symlink")
|
||||
self.api.db.link_image_to_prompt(prompt, str(spelled / "b.png"))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertIn("1 database link moves", plan["details"][0]["detail"])
|
||||
await self.apply_plan(plan)
|
||||
[image] = self.api.db.get_prompt_images(prompt)
|
||||
self.assertEqual(
|
||||
_locator(image["image_path"]), _locator(self.output_dir / "a.png")
|
||||
)
|
||||
self.assertEqual(self.linked_files(), [])
|
||||
|
||||
|
||||
class TestTagsAndReview(SafetyTestCase):
|
||||
"""I5: tags join the kept copy; pending review items of the copy settle."""
|
||||
|
||||
def tag(self, path, tags):
|
||||
data = Path(path).read_bytes()
|
||||
info = os.stat(path)
|
||||
return self.api.db.replace_image_tags(
|
||||
locator=_locator(path),
|
||||
sha256=hashlib.sha256(data).hexdigest(),
|
||||
size=info.st_size,
|
||||
mtime_ns=info.st_mtime_ns,
|
||||
expected_epoch=self.api.db.get_execution_epoch(),
|
||||
expected_tag_revision=0,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def tags_of(self, path):
|
||||
data = Path(path).read_bytes()
|
||||
return self.api.db.read_image_tags(
|
||||
_locator(path), hashlib.sha256(data).hexdigest()
|
||||
)["tags"]
|
||||
|
||||
def queue(self, path):
|
||||
data = Path(path).read_bytes()
|
||||
info = os.stat(path)
|
||||
return self.api.db.record_review_item(
|
||||
locator=_locator(path),
|
||||
spelled_path=os.path.realpath(str(path)),
|
||||
sha256=hashlib.sha256(data).hexdigest(),
|
||||
size=info.st_size,
|
||||
mtime_ns=info.st_mtime_ns,
|
||||
epoch=self.api.db.get_execution_epoch(),
|
||||
resolver_version=1,
|
||||
evidence={"candidates": []},
|
||||
candidate_count=0,
|
||||
limited=True,
|
||||
)
|
||||
|
||||
def pending(self):
|
||||
return self.api.db.review_queue_page(limit=50)["items"]
|
||||
|
||||
async def test_tags_move_as_a_union_and_pending_items_settle(self):
|
||||
self.write("a.png", b"same", mtime=1)
|
||||
self.write("b.png", b"same", mtime=2)
|
||||
keeper, removed = self.output_dir / "a.png", self.output_dir / "b.png"
|
||||
self.tag(keeper, ["kept"])
|
||||
self.tag(removed, ["moved", "kept"])
|
||||
self.assertTrue(self.queue(removed))
|
||||
job, group = await self.found_group()
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertIn("image tags join the kept copy", plan["details"][0]["detail"])
|
||||
self.assertIn("1 Needs review item is settled", plan["details"][0]["detail"])
|
||||
await self.apply_plan(plan)
|
||||
self.assertEqual(self.tags_of(keeper), ["kept", "moved"])
|
||||
self.assertEqual(self.pending(), [])
|
||||
|
||||
|
||||
class TestPreviewDoesNotHash(SafetyTestCase):
|
||||
"""M9: the preview binds by size, mtime and inode; the job hashes."""
|
||||
|
||||
async def test_preview_reads_no_file_and_spots_a_touched_one(self):
|
||||
self.trio()
|
||||
job, group = await self.found_group()
|
||||
calls = []
|
||||
real = self.api._calculate_file_hash
|
||||
|
||||
def spy(path):
|
||||
calls.append(path)
|
||||
return real(path)
|
||||
|
||||
self.api._calculate_file_hash = spy
|
||||
os.utime(self.output_dir / "b.png", (5_000_000, 5_000_000))
|
||||
plan = await self.review(job.id, [self.keep_of(group, "a.png")])
|
||||
self.assertEqual(calls, [])
|
||||
self.assertEqual(plan["scope"]["count"], 1)
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"], ["1 file changed since the search"]
|
||||
)
|
||||
cleanup = await self.apply_plan(plan)
|
||||
self.assertEqual(cleanup.counters["removed"], 1)
|
||||
self.assertGreaterEqual(len(calls), 2) # the job hashes keeper and file
|
||||
@@ -588,11 +588,22 @@ class TestSafetyBackup(ApplyTestCase):
|
||||
async def test_a_failed_replace_leaves_no_file(self):
|
||||
from py.api import maintenance
|
||||
|
||||
with patch.object(maintenance.os, "replace", side_effect=OSError("busy")):
|
||||
# A filesystem without hard links falls back to reserve-and-replace.
|
||||
with (
|
||||
patch.object(maintenance.os, "link", side_effect=OSError("no hard links")),
|
||||
patch.object(maintenance.os, "replace", side_effect=OSError("busy")),
|
||||
):
|
||||
with self.assertRaises(maintenance.SafetyBackupFailed):
|
||||
maintenance.take_safety_backup(self.api.db.model)
|
||||
self.assertEqual(self.backups(), [])
|
||||
self.assertEqual(self.strays(), [])
|
||||
with patch.object(maintenance.os, "link", side_effect=FileExistsError):
|
||||
with self.assertRaises(maintenance.SafetyBackupFailed):
|
||||
maintenance.take_safety_backup(self.api.db.model)
|
||||
with patch.object(maintenance.os, "link", side_effect=OSError("no hard links")):
|
||||
name = maintenance.take_safety_backup(self.api.db.model)
|
||||
self.assertEqual(self.backups(), [name])
|
||||
self.assertEqual(self.strays(), [])
|
||||
|
||||
async def test_a_backup_that_cannot_be_finished_leaves_no_file(self):
|
||||
from py.api import maintenance
|
||||
@@ -689,7 +700,11 @@ class TestBackupPoolAndDiskSpace(ApplyTestCase):
|
||||
from py.api import maintenance
|
||||
|
||||
self.save("Orphan")
|
||||
with patch.object(maintenance, "_free_bytes", return_value=10 * 2**30):
|
||||
# SQLite's temporary folder on the database's disk: the needs add up.
|
||||
with (
|
||||
patch.object(maintenance, "_free_bytes", return_value=10 * 2**30),
|
||||
patch.object(maintenance, "sqlite_temp_dir", return_value=self.tmpdir),
|
||||
):
|
||||
plan = await self.preview("unused_components")
|
||||
optimize = await self.preview("optimize")
|
||||
size = self.database_bytes()
|
||||
@@ -846,3 +861,207 @@ if __name__ == "__main__": # pragma: no cover
|
||||
import unittest
|
||||
|
||||
unittest.main()
|
||||
|
||||
|
||||
class TestFinalReviewFixes(ApplyTestCase):
|
||||
"""PR #313 final review: I3, M1-M5."""
|
||||
|
||||
def duplicate_of(self, kept_text, copy_text):
|
||||
kept = self.save(kept_text)
|
||||
with self.api.db.model.get_connection() as conn:
|
||||
copy = conn.execute(
|
||||
"INSERT INTO prompts(text, created_at) VALUES(?, '2999-01-01')",
|
||||
(copy_text,),
|
||||
).lastrowid
|
||||
return kept, copy
|
||||
|
||||
def test_sqlite_temp_dir_follows_sqlite_search_order(self):
|
||||
from py.api.maintenance import sqlite_temp_dir
|
||||
|
||||
def usable(*good):
|
||||
return lambda path: path in good
|
||||
|
||||
posix = {"windows": False}
|
||||
self.assertEqual(
|
||||
sqlite_temp_dir({}, usable=usable("/var/tmp", "/tmp"), **posix), "/var/tmp"
|
||||
)
|
||||
self.assertEqual(
|
||||
sqlite_temp_dir({}, usable=usable("/usr/tmp", "/tmp"), **posix), "/usr/tmp"
|
||||
)
|
||||
env = {"SQLITE_TMPDIR": "/fast", "TMPDIR": "/t"}
|
||||
self.assertEqual(
|
||||
sqlite_temp_dir(env, usable=usable("/fast", "/t", "/tmp"), **posix),
|
||||
"/fast",
|
||||
)
|
||||
self.assertEqual(
|
||||
sqlite_temp_dir(env, usable=usable("/t", "/tmp"), **posix), "/t"
|
||||
)
|
||||
self.assertEqual(
|
||||
sqlite_temp_dir({}, usable=usable(), **posix), os.path.abspath(".")
|
||||
)
|
||||
windows = {"windows": True, "usable": usable()}
|
||||
env = {"TMP": "C:\\a", "TEMP": "C:\\b", "USERPROFILE": "C:\\u"}
|
||||
self.assertEqual(sqlite_temp_dir(env, **windows), "C:\\a")
|
||||
self.assertEqual(sqlite_temp_dir({"TEMP": "C:\\b"}, **windows), "C:\\b")
|
||||
self.assertEqual(sqlite_temp_dir({"USERPROFILE": "C:\\u"}, **windows), "C:\\u")
|
||||
|
||||
async def test_optimize_measures_sqlites_temp_folder(self):
|
||||
from py.api import maintenance
|
||||
|
||||
measured = []
|
||||
|
||||
def free(path):
|
||||
measured.append(path)
|
||||
return 10 * 2**30
|
||||
|
||||
other = os.path.join(self.tmpdir, "sqlite-temp")
|
||||
os.mkdir(other)
|
||||
with (
|
||||
patch.object(maintenance, "sqlite_temp_dir", return_value=other),
|
||||
patch.object(maintenance, "_free_bytes", side_effect=free),
|
||||
patch.object(
|
||||
maintenance, "_device", side_effect=lambda path: path # another disk
|
||||
),
|
||||
):
|
||||
plan = await self.preview("optimize")
|
||||
self.assertIn(other, measured)
|
||||
size = self.database_bytes()
|
||||
self.assertEqual(
|
||||
plan["disk"]["required_bytes"], 2 * size + maintenance.DISK_MARGIN_BYTES
|
||||
)
|
||||
|
||||
async def test_a_plan_used_during_the_disk_check_is_plan_used(self):
|
||||
from py.api import maintenance
|
||||
|
||||
self.save("Orphan")
|
||||
plan = await self.preview("unused_components")
|
||||
real = maintenance.disk_plan
|
||||
|
||||
def racing(*args, **kwargs):
|
||||
# A parallel apply consumed the token while this one measured.
|
||||
self.api._maintenance_plans.consume(
|
||||
plan["token"], "unused_components", job_id="f" * 32
|
||||
)
|
||||
return real(*args, **kwargs)
|
||||
|
||||
with patch.object(maintenance, "disk_plan", side_effect=racing):
|
||||
reply = await self.apply(plan["token"], 409)
|
||||
self.assertEqual(reply["code"], "plan-used")
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
|
||||
async def test_the_backup_decision_is_frozen_when_the_job_starts(self):
|
||||
PromptManagerConfig.SAFETY_BACKUP = False
|
||||
self.save("Orphan")
|
||||
plan = await self.preview("unused_components")
|
||||
real = self.api._jobs.start
|
||||
|
||||
def start(*args, **kwargs):
|
||||
job = real(*args, **kwargs)
|
||||
PromptManagerConfig.SAFETY_BACKUP = True # switched on meanwhile
|
||||
return job
|
||||
|
||||
with patch.object(self.api._jobs, "start", side_effect=start):
|
||||
body = await self.apply(plan["token"])
|
||||
job = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
self.assertEqual(self.backups(), [])
|
||||
self.assertTrue(job.status_line.endswith("No safety backup"))
|
||||
|
||||
async def test_startup_sweeps_interrupted_backup_copies(self):
|
||||
stray = os.path.join(self.tmpdir, ".pm-safety-crashed.tmp")
|
||||
with open(stray, "wb") as handle:
|
||||
handle.write(b"half a backup")
|
||||
PromptManagerAPI(db=self.api.db)
|
||||
self.assertFalse(os.path.exists(stray))
|
||||
|
||||
async def test_a_backup_is_linked_into_place_without_a_reserved_name(self):
|
||||
from py.api import maintenance
|
||||
|
||||
linked = []
|
||||
real = os.link
|
||||
|
||||
def link(source, target):
|
||||
linked.append(target)
|
||||
return real(source, target)
|
||||
|
||||
with (
|
||||
patch.object(maintenance.os, "link", side_effect=link),
|
||||
patch.object(
|
||||
maintenance, "_reserve_backup_name", side_effect=AssertionError
|
||||
),
|
||||
):
|
||||
name = maintenance.take_safety_backup(self.api.db.model)
|
||||
self.assertEqual([os.path.basename(p) for p in linked], [name])
|
||||
self.assertEqual(self.strays(), [])
|
||||
|
||||
async def test_optimize_runs_alone_and_says_the_database_is_locked(self):
|
||||
from py.api.jobs import JobBusy
|
||||
|
||||
plan = await self.preview("optimize")
|
||||
self.assertIn(
|
||||
"Runs alone: the database is locked while it runs",
|
||||
" ".join(row["detail"] for row in plan["details"]),
|
||||
)
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def held(job):
|
||||
await gate.wait()
|
||||
|
||||
read = self.api._jobs.start(
|
||||
"health", "Health check", admission="read", runner=held
|
||||
)
|
||||
busy = await self.apply(plan["token"], 409)
|
||||
self.assertEqual(busy["running"]["id"], read.id)
|
||||
gate.set()
|
||||
await asyncio.wait_for(asyncio.shield(read.task), WAIT)
|
||||
gate.clear()
|
||||
body = await self.apply(plan["token"])
|
||||
self.assertEqual(body["job"]["admission"], "exclusive")
|
||||
with self.assertRaises(JobBusy):
|
||||
self.api._jobs.start(
|
||||
"health", "Health check", admission="read", runner=held
|
||||
)
|
||||
job = self.api._jobs.get(body["job"]["id"])
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
|
||||
async def test_a_link_collision_is_protected_not_swallowed(self):
|
||||
kept, copy = self.duplicate_of("A cat", "A CAT")
|
||||
path = os.path.join(str(self.output_dir), "both.png")
|
||||
self.api.db.link_image_to_prompt(kept, path)
|
||||
self.api.db.link_image_to_prompt(copy, path)
|
||||
plan, _summary, job = await self.run_apply("duplicate_prompts")
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"], ["1 group links one image twice"]
|
||||
)
|
||||
self.assertIsNotNone(self.api.db.get_prompt_by_id(copy))
|
||||
self.assertEqual(len(self.api.db.get_prompt_images(copy)), 1)
|
||||
|
||||
async def test_a_duplicate_with_a_library_title_is_protected(self):
|
||||
kept, copy = self.duplicate_of("A dog", "A DOG")
|
||||
with self.api.db.model.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
||||
(copy, "My dog recipe"),
|
||||
)
|
||||
plan, _summary, _job = await self.run_apply("duplicate_prompts")
|
||||
self.assertEqual(plan["scope"]["count"], 0)
|
||||
self.assertEqual(
|
||||
plan["protected"]["reasons"], ["1 group has its own Library title"]
|
||||
)
|
||||
self.assertIsNotNone(self.api.db.get_prompt_by_id(copy))
|
||||
|
||||
def test_the_strict_merge_raises_instead_of_dropping_a_link(self):
|
||||
from database.maintenance import MergeConflict
|
||||
|
||||
kept, copy = self.duplicate_of("A bird", "A BIRD")
|
||||
path = os.path.join(str(self.output_dir), "bird.png")
|
||||
self.api.db.link_image_to_prompt(kept, path)
|
||||
self.api.db.link_image_to_prompt(copy, path)
|
||||
with self.assertRaises(MergeConflict):
|
||||
with self.api.db._pair_transaction() as conn:
|
||||
self.api.db._merge_duplicate_group(conn, [kept, copy], strict=True)
|
||||
self.assertEqual(len(self.api.db.get_prompt_images(copy)), 1)
|
||||
# The legacy cleanup keeps its old behaviour (it logs and goes on).
|
||||
with self.assertLogs("prompt_manager.database", "ERROR"):
|
||||
self.assertEqual(self.api.db.cleanup_duplicates(), 1)
|
||||
|
||||
@@ -63,7 +63,7 @@ class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):
|
||||
)
|
||||
|
||||
async def test_admission_names_and_explicit_read(self):
|
||||
self.assertEqual(ADMISSIONS, ("read", "guarded", "records"))
|
||||
self.assertEqual(ADMISSIONS, ("read", "guarded", "records", "exclusive"))
|
||||
registry = JobRegistry()
|
||||
job = registry.start(
|
||||
"report", "Report", admission="read", scope={}, runner=_noop
|
||||
@@ -72,7 +72,7 @@ class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertEqual(job.admission, "read")
|
||||
self.assertFalse(job.record_changing)
|
||||
with self.assertRaises(ValueError):
|
||||
registry.start("x", "X", admission="exclusive", scope={}, runner=_noop)
|
||||
registry.start("x", "X", admission="bogus", scope={}, runner=_noop)
|
||||
with self.assertRaises(ValueError):
|
||||
registry.start(
|
||||
"x",
|
||||
|
||||
@@ -398,7 +398,7 @@ class TestPreviews(MaintenanceDatabaseCase):
|
||||
self.assertEqual(draft.scope, {"count": 1, "label": "The whole database"})
|
||||
self.assertEqual(
|
||||
[row["label"] for row in draft.details],
|
||||
["Database file", "Free pages", "Free space needed"],
|
||||
["Database file", "Free pages", "Lock", "Free space needed"],
|
||||
)
|
||||
self.assertNotIn(str(self.root), str(draft.public()))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user