447 lines
18 KiB
Python
447 lines
18 KiB
Python
"""#443 (D152): the Library delete, archive and restore routes, and ingestion after."""
|
|
|
|
import json
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
from tests.test_api_admin import AdminAPITestCase # first: it stubs ComfyUI
|
|
from database.execution_history_schema import rotate_history_epoch # noqa: E402
|
|
from database.prompt_archive import TABLE as ARCHIVE # noqa: E402
|
|
from database.safety_backup import SafetyBackupFailed # noqa: E402
|
|
from py.config import PromptManagerConfig # noqa: E402
|
|
from utils.hashing import generate_prompt_hash # noqa: E402
|
|
|
|
PREVIEW = "/prompt_manager/library/delete/preview"
|
|
DELETE = "/prompt_manager/library/delete"
|
|
RESTORE = "/prompt_manager/library/restore"
|
|
|
|
|
|
def graph(positive, negative="blur"):
|
|
return {
|
|
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}},
|
|
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
|
},
|
|
}
|
|
|
|
|
|
class LibraryDeleteRouteCase(AdminAPITestCase):
|
|
def epoch(self):
|
|
return self.api.db.get_library_options()["library_epoch"]
|
|
|
|
def save(self, text, **metadata):
|
|
return self.api.db.save_prompt(
|
|
text, prompt_hash=generate_prompt_hash(text), **metadata
|
|
)
|
|
|
|
def paired(self, positive, negative="blur"):
|
|
pair = self.api.db.save_prompt_pair(positive, negative)
|
|
return self.api.db.get_prompt_pair(pair)["positive_prompt_id"]
|
|
|
|
def rows(self, sql, params=()):
|
|
with self.api.db.model.get_connection() as conn:
|
|
return [tuple(row) for row in conn.execute(sql, params).fetchall()]
|
|
|
|
def ids(self):
|
|
return sorted(row[0] for row in self.rows("SELECT id FROM prompts"))
|
|
|
|
def archived(self):
|
|
return sorted(row[0] for row in self.rows(f"SELECT prompt_id FROM {ARCHIVE}"))
|
|
|
|
def backups(self):
|
|
folder = os.path.dirname(self.api.db.model.db_path)
|
|
return sorted(name for name in os.listdir(folder) if ".backup_" in name)
|
|
|
|
async def post(self, url, body):
|
|
response = await self.client.request("POST", url, json=body)
|
|
return response.status, await response.json()
|
|
|
|
async def review(self, ids):
|
|
status, body = await self.post(
|
|
PREVIEW, {"prompt_ids": ids, "expected_epoch": self.epoch()}
|
|
)
|
|
self.assertEqual(status, 200, body)
|
|
return body
|
|
|
|
async def confirm(self, review, backup=None):
|
|
return await self.post(
|
|
DELETE,
|
|
{
|
|
"archive_ids": review["archive"],
|
|
"delete_ids": review["delete"],
|
|
"expected_epoch": review["library_epoch"],
|
|
"backup": backup or review["backup"],
|
|
},
|
|
)
|
|
|
|
def policy(self, on):
|
|
previous = PromptManagerConfig.SAFETY_BACKUP
|
|
PromptManagerConfig.SAFETY_BACKUP = on
|
|
self.addCleanup(setattr, PromptManagerConfig, "SAFETY_BACKUP", previous)
|
|
|
|
|
|
class TestLibraryDeleteRoutes(LibraryDeleteRouteCase):
|
|
async def test_preview_then_delete_archives_deletes_and_backs_up(self):
|
|
self.policy(True)
|
|
plain = self.save("Plain")
|
|
guarded = self.save("Guarded", tags=["__protected__"])
|
|
positive = self.paired("Cat")
|
|
review = await self.review([plain, guarded, positive, 99999])
|
|
self.assertEqual(
|
|
{
|
|
key: review[key]
|
|
for key in ("delete", "archive", "protected", "not_found")
|
|
},
|
|
{
|
|
"delete": [plain],
|
|
"archive": [positive],
|
|
"protected": [guarded],
|
|
"not_found": 1,
|
|
},
|
|
)
|
|
self.assertEqual(review["backup"], "on")
|
|
self.assertTrue(review["disk"]["sufficient"])
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(
|
|
(body["deleted"], body["archived"], body["changed"]), (1, 1, 0)
|
|
)
|
|
self.assertEqual(self.backups(), [body["backup"]])
|
|
self.assertNotIn(os.sep, body["backup"])
|
|
self.assertNotIn(plain, self.ids())
|
|
self.assertEqual(self.archived(), [positive])
|
|
# The backup still holds the deleted prompt.
|
|
import sqlite3
|
|
|
|
path = os.path.join(os.path.dirname(self.api.db.model.db_path), body["backup"])
|
|
with sqlite3.connect(path) as copy:
|
|
self.assertEqual(
|
|
copy.execute(
|
|
"SELECT COUNT(*) FROM prompts WHERE id=?", (plain,)
|
|
).fetchone()[0],
|
|
1,
|
|
)
|
|
copy.close()
|
|
# The hold ended: nothing is listed as running.
|
|
self.assertEqual(self.api._jobs.running(), [])
|
|
|
|
async def test_archive_only_takes_no_backup(self):
|
|
self.policy(True)
|
|
positive = self.paired("Cat")
|
|
shared = self.save("Shared")
|
|
self.api.db.save_prompt_pair("Dog", "Shared") # archive-only (#482)
|
|
self.api.db.link_image_to_prompt(shared, "/tmp/shared-443.png")
|
|
# Nothing could be deleted: no disk plan. (#482: a deletable prompt has one.)
|
|
self.assertIsNone((await self.review([shared]))["disk"])
|
|
review = await self.review([positive])
|
|
self.assertIsNotNone(review["disk"])
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual((status, body["archived"], body["backup"]), (200, 1, None))
|
|
self.assertEqual(self.backups(), [])
|
|
|
|
async def test_policy_off_takes_none_unless_the_review_promised_one(self):
|
|
plain = self.save("Plain")
|
|
other = self.save("Other")
|
|
self.policy(False)
|
|
review = await self.review([plain])
|
|
self.assertEqual((review["backup"], review["disk"]), ("off", None))
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual((status, body["deleted"], body["backup"]), (200, 1, None))
|
|
# Reviewed with the policy on, turned off before the confirm: kept.
|
|
review = await self.review([other])
|
|
status, body = await self.confirm(review, backup="on")
|
|
self.assertEqual(status, 200)
|
|
self.assertIsNotNone(body["backup"])
|
|
|
|
async def test_a_failed_backup_changes_nothing(self):
|
|
self.policy(True)
|
|
plain = self.save("Plain")
|
|
positive = self.paired("Cat")
|
|
review = await self.review([plain, positive])
|
|
with patch(
|
|
"py.api.maintenance.take_safety_backup",
|
|
side_effect=SafetyBackupFailed("x"),
|
|
):
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual((status, body["code"]), (500, "backup-failed"))
|
|
self.assertIn("nothing was changed", body["error"])
|
|
self.assertIn(plain, self.ids())
|
|
self.assertEqual(self.archived(), [])
|
|
|
|
async def test_a_short_disk_refuses_before_any_change(self):
|
|
self.policy(True)
|
|
plain = self.save("Plain")
|
|
from py.api import maintenance
|
|
|
|
with patch.object(maintenance, "_free_bytes", return_value=0):
|
|
review = await self.review([plain])
|
|
self.assertFalse(review["disk"]["sufficient"])
|
|
self.assertIn("Nothing was changed", review["disk"]["message"])
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual((status, body["code"]), (507, "insufficient-space"))
|
|
self.assertIn(plain, self.ids())
|
|
|
|
async def test_busy_beside_a_record_changing_job(self):
|
|
plain = self.save("Plain")
|
|
review = await self.review([plain])
|
|
async with self.api._jobs.hold("scan", "Scan", admission="records"):
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual((status, body["code"]), (409, "busy"))
|
|
self.assertEqual(body["running"]["title"], "Scan")
|
|
self.assertIn(plain, self.ids())
|
|
|
|
async def test_restored_storage_answers_409_on_every_route(self):
|
|
plain = self.save("Plain")
|
|
review = await self.review([plain])
|
|
rotate_history_epoch(self.api.db.model.get_connection())
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual((status, body["code"]), (409, "storage-changed"))
|
|
for url in (PREVIEW, RESTORE):
|
|
status, body = await self.post(
|
|
url, {"prompt_ids": [plain], "expected_epoch": review["library_epoch"]}
|
|
)
|
|
self.assertEqual((status, body["code"]), (409, "storage-changed"))
|
|
self.assertIn(plain, self.ids())
|
|
|
|
async def test_restore_route(self):
|
|
positive = self.paired("Cat")
|
|
await self.confirm(await self.review([positive]))
|
|
status, body = await self.post(
|
|
RESTORE, {"prompt_ids": [positive, 99999], "expected_epoch": self.epoch()}
|
|
)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual((body["restored"], body["not_found"]), (1, 1))
|
|
self.assertEqual(self.archived(), [])
|
|
|
|
async def test_bodies_are_validated(self):
|
|
plain = self.save("Plain")
|
|
epoch = self.epoch()
|
|
selection = {"prompt_ids": [plain], "expected_epoch": epoch}
|
|
for url, body in (
|
|
(PREVIEW, {**selection, "extra": 1}),
|
|
(PREVIEW, {"prompt_ids": [], "expected_epoch": epoch}),
|
|
(PREVIEW, {"prompt_ids": [True], "expected_epoch": epoch}),
|
|
(PREVIEW, {"prompt_ids": list(range(1, 1002)), "expected_epoch": epoch}),
|
|
(PREVIEW, {"prompt_ids": [plain], "expected_epoch": "nope"}),
|
|
(RESTORE, {"prompt_ids": [plain]}),
|
|
(DELETE, {**selection}),
|
|
(
|
|
DELETE,
|
|
{
|
|
"archive_ids": [],
|
|
"delete_ids": [],
|
|
"expected_epoch": epoch,
|
|
"backup": "on",
|
|
},
|
|
),
|
|
(
|
|
DELETE,
|
|
{
|
|
"archive_ids": [plain],
|
|
"delete_ids": [plain],
|
|
"expected_epoch": epoch,
|
|
"backup": "on",
|
|
},
|
|
),
|
|
(
|
|
DELETE,
|
|
{
|
|
"archive_ids": [],
|
|
"delete_ids": [plain],
|
|
"expected_epoch": epoch,
|
|
"backup": "maybe",
|
|
},
|
|
),
|
|
(
|
|
DELETE,
|
|
{
|
|
"archive_ids": list(range(1, 600)),
|
|
"delete_ids": list(range(600, 1100)),
|
|
"expected_epoch": epoch,
|
|
"backup": "off",
|
|
},
|
|
),
|
|
):
|
|
with self.subTest(url=url, body=body):
|
|
status, answer = await self.post(url, body)
|
|
self.assertEqual(status, 400, answer)
|
|
self.assertFalse(answer["success"])
|
|
for url in (PREVIEW, DELETE, RESTORE):
|
|
response = await self.client.request(
|
|
"POST", url, data="[1]", headers={"Content-Type": "application/json"}
|
|
)
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn(plain, self.ids())
|
|
|
|
async def test_the_legacy_bulk_delete_is_unchanged(self):
|
|
plain = self.save("Plain")
|
|
positive = self.paired("Cat")
|
|
response = await self.client.request(
|
|
"POST", "/prompt_manager/bulk/delete", json={"prompt_ids": [positive]}
|
|
)
|
|
self.assertEqual(response.status, 409)
|
|
response = await self.client.request(
|
|
"POST", "/prompt_manager/bulk/delete", json={"prompt_ids": [plain]}
|
|
)
|
|
self.assertEqual((await response.json())["deleted_count"], 1)
|
|
self.assertEqual(self.archived(), [])
|
|
|
|
async def test_shelf_parameter_and_cursor_scope(self):
|
|
kept = [self.paired(f"Kept {i}") for i in range(3)]
|
|
gone = [self.paired(f"Gone {i}") for i in range(3)]
|
|
await self.confirm(await self.review(gone))
|
|
response = await self.client.request(
|
|
"GET", "/prompt_manager/library?limit=2&shelf=active"
|
|
)
|
|
page = await response.json()
|
|
self.assertEqual(page["archived_total"], 3)
|
|
self.assertTrue(set(item["id"] for item in page["items"]) <= set(kept))
|
|
response = await self.client.request(
|
|
"GET",
|
|
"/prompt_manager/library?limit=2&shelf=archived"
|
|
f"&cursor={page['next_cursor']}",
|
|
)
|
|
self.assertEqual(response.status, 400)
|
|
response = await self.client.request(
|
|
"GET", "/prompt_manager/library?shelf=archived"
|
|
)
|
|
archived = await response.json()
|
|
self.assertEqual(sorted(item["id"] for item in archived["items"]), gone)
|
|
self.assertTrue(all(item["archived"] for item in archived["items"]))
|
|
response = await self.client.request("GET", "/prompt_manager/library?shelf=all")
|
|
self.assertEqual(response.status, 400)
|
|
|
|
|
|
class TestPairedDeleteRoutes(LibraryDeleteRouteCase):
|
|
"""#482 (D160): Delete permanently with a prompt's own saved pairs."""
|
|
|
|
async def confirm_paired(self, review, ids, backup=None, digests=None):
|
|
return await self.post(
|
|
DELETE,
|
|
{
|
|
"archive_ids": review["archive"],
|
|
"delete_ids": review["delete"],
|
|
"expected_epoch": review["library_epoch"],
|
|
"backup": backup or review["backup"],
|
|
"delete_paired_ids": ids,
|
|
"paired_digests": digests
|
|
or {str(pid): review["review_digests"][str(pid)] for pid in ids},
|
|
},
|
|
)
|
|
|
|
async def test_preview_reports_deletable_and_a_paired_only_delete_backs_up(self):
|
|
self.policy(True)
|
|
cat = self.paired("Cat")
|
|
shared = self.save("Shared")
|
|
self.api.db.save_prompt_pair("Dog", "Shared")
|
|
self.api.db.link_image_to_prompt(shared, "/tmp/shared-482.png")
|
|
review = await self.review([cat, shared])
|
|
self.assertEqual(review["archive"], [cat, shared])
|
|
self.assertEqual(review["deletable"], [cat])
|
|
self.assertEqual(list(review["review_digests"]), [str(cat)])
|
|
self.assertEqual(
|
|
review["paired_counts"],
|
|
{"pairs": 1, "archived_pairs": 0, "variants": 1, "kept_negatives": 1},
|
|
)
|
|
self.assertEqual(review["paired_image_links"], 0)
|
|
# A selection with nothing in the delete class still checks the disk.
|
|
self.assertTrue(review["disk"]["sufficient"])
|
|
status, body = await self.confirm_paired(review, [cat])
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(
|
|
(body["deleted"], body["archived"], body["pairs_removed"]), (1, 1, 1)
|
|
)
|
|
self.assertEqual(self.backups(), [body["backup"]])
|
|
self.assertNotIn(cat, self.ids())
|
|
self.assertEqual(self.archived(), [shared])
|
|
|
|
async def test_a_body_without_the_new_keys_is_the_443_delete(self):
|
|
cat = self.paired("Cat")
|
|
review = await self.review([cat])
|
|
status, body = await self.confirm(review)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual((body["deleted"], body["archived"]), (0, 1))
|
|
self.assertIn(cat, self.ids())
|
|
|
|
async def test_paired_bodies_are_validated(self):
|
|
cat = self.paired("Cat")
|
|
plain = self.save("Plain")
|
|
review = await self.review([cat, plain])
|
|
digest = review["review_digests"][str(cat)]
|
|
base = {
|
|
"archive_ids": review["archive"],
|
|
"delete_ids": review["delete"],
|
|
"expected_epoch": review["library_epoch"],
|
|
"backup": "off",
|
|
}
|
|
cases = [
|
|
{"delete_paired_ids": [cat]}, # without its digests
|
|
{"delete_paired_ids": [plain], "paired_digests": {str(plain): digest}},
|
|
{"delete_paired_ids": [cat], "paired_digests": {}},
|
|
{"delete_paired_ids": [cat], "paired_digests": {str(cat): "x"}},
|
|
{"delete_paired_ids": "all", "paired_digests": {}},
|
|
]
|
|
for extra in cases:
|
|
with self.subTest(extra=extra):
|
|
status, body = await self.post(DELETE, {**base, **extra})
|
|
self.assertEqual((status, body["code"]), (400, "invalid-request"))
|
|
self.assertIn(cat, self.ids())
|
|
|
|
async def test_a_short_disk_refuses_a_paired_only_delete(self):
|
|
self.policy(True)
|
|
cat = self.paired("Cat")
|
|
from py.api import maintenance
|
|
|
|
with patch.object(maintenance, "_free_bytes", return_value=0):
|
|
review = await self.review([cat])
|
|
self.assertFalse(review["disk"]["sufficient"])
|
|
status, body = await self.confirm_paired(review, [cat])
|
|
self.assertEqual((status, body["code"]), (507, "insufficient-space"))
|
|
self.assertIn(cat, self.ids())
|
|
|
|
|
|
class TestIngestionAfterArchive(LibraryDeleteRouteCase):
|
|
def scan(self, name, positive, negative="blur"):
|
|
meta = {"prompt": json.dumps(graph(positive, negative))}
|
|
path = self.output_dir / name
|
|
path.write_bytes(b"png")
|
|
return self.api._ingest_scanned_file(path, meta)
|
|
|
|
async def test_a_scan_after_archive_creates_no_duplicate_and_keeps_it_archived(
|
|
self,
|
|
):
|
|
self.assertEqual(self.scan("first.png", "Sunset"), "added")
|
|
positive = self.api.db.get_prompt_by_hash(generate_prompt_hash("Sunset"))["id"]
|
|
await self.confirm(await self.review([positive]))
|
|
before = self.ids()
|
|
self.scan("second.png", "Sunset")
|
|
self.scan("third.png", " sunset ", "other negative")
|
|
self.assertEqual(
|
|
[i for i in self.ids() if i not in before],
|
|
[
|
|
self.api.db.get_prompt_by_hash(generate_prompt_hash("other negative"))[
|
|
"id"
|
|
]
|
|
],
|
|
)
|
|
self.assertEqual(self.archived(), [positive])
|
|
linked = self.rows(
|
|
"SELECT COUNT(*) FROM generated_images WHERE prompt_id=?", (positive,)
|
|
)
|
|
self.assertEqual(linked, [(3,)])
|
|
response = await self.client.request("GET", "/prompt_manager/library")
|
|
self.assertNotIn(
|
|
positive, [item["id"] for item in (await response.json())["items"]]
|
|
)
|
|
|
|
async def test_a_scan_after_delete_creates_the_prompt_again(self):
|
|
plain = self.save("Lonely sunset")
|
|
await self.confirm(await self.review([plain]))
|
|
self.assertEqual(self.scan("again.png", "Lonely sunset"), "added")
|
|
again = self.api.db.get_prompt_by_hash(generate_prompt_hash("Lonely sunset"))
|
|
self.assertIsNotNone(again)
|
|
self.assertNotEqual(again["id"], plain)
|
|
self.assertEqual(self.archived(), [])
|