Vito Sansevero
507c8a8993
Each owned test that opens a PromptDatabase now calls close_all() before removing the file, and temp directories holding a database use ignore_cleanup_errors, so the suite tears down cleanly on Windows. The image API fixture lays out a fake ComfyUI tree (<tmp>/output) and pins GalleryConfig.path_anchors to it so responses can be checked for root-relative paths.
240 lines
8.5 KiB
Python
240 lines
8.5 KiB
Python
"""Route tests for GET /prompt_manager/backup and POST /prompt_manager/restore.
|
|
|
|
The backup must include rows still sitting in the WAL, the restore must
|
|
stream the upload to disk, reject anything that is not a healthy prompts
|
|
database and leave the live database untouched when it does.
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import aiohttp
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import AioHTTPTestCase
|
|
|
|
from database.operations import PromptDatabase
|
|
from py.api import PromptManagerAPI
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
def _texts(path):
|
|
conn = sqlite3.connect(path)
|
|
try:
|
|
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
class BackupRestoreTestCase(AioHTTPTestCase):
|
|
async def get_application(self):
|
|
self.tmpdir = tempfile.mkdtemp()
|
|
self.db_path = os.path.join(self.tmpdir, "prompts.db")
|
|
|
|
app = web.Application()
|
|
routes = web.RouteTableDef()
|
|
self.api = PromptManagerAPI()
|
|
self.api.db = PromptDatabase(self.db_path)
|
|
self.api.add_routes(routes)
|
|
app.router.add_routes(routes)
|
|
return app
|
|
|
|
async def tearDownAsync(self):
|
|
self.api.db.close_all()
|
|
for name in os.listdir(self.tmpdir):
|
|
os.unlink(os.path.join(self.tmpdir, name))
|
|
os.rmdir(self.tmpdir)
|
|
|
|
def _save(self, text):
|
|
return self.api.db.save_prompt(
|
|
text=text, prompt_hash=generate_prompt_hash(text)
|
|
)
|
|
|
|
def _write(self, name, data):
|
|
path = os.path.join(self.tmpdir, name)
|
|
with open(path, "wb") as fh:
|
|
fh.write(data)
|
|
return path
|
|
|
|
def _make_db_file(self, texts, name="upload.db"):
|
|
path = os.path.join(self.tmpdir, name)
|
|
db = PromptDatabase(path)
|
|
for text in texts:
|
|
db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
|
db.close_all()
|
|
return path
|
|
|
|
async def _restore(self, path, field="database_file"):
|
|
with open(path, "rb") as fh:
|
|
form = aiohttp.FormData()
|
|
form.add_field(
|
|
field,
|
|
fh.read(),
|
|
filename=os.path.basename(path),
|
|
content_type="application/octet-stream",
|
|
)
|
|
return await self.client.post("/prompt_manager/restore", data=form)
|
|
|
|
|
|
class TestBackupRoute(BackupRestoreTestCase):
|
|
async def test_backup_download_contains_rows_still_in_the_wal(self):
|
|
for n in range(3):
|
|
self._save(f"backup me {n}")
|
|
self.assertGreater(os.path.getsize(self.db_path + "-wal"), 0)
|
|
created = []
|
|
real_mkstemp = tempfile.mkstemp
|
|
|
|
def spy_mkstemp(*args, **kwargs):
|
|
result = real_mkstemp(*args, **kwargs)
|
|
created.append(result[1])
|
|
return result
|
|
|
|
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
|
|
resp = await self.client.get("/prompt_manager/backup")
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.content_type, "application/octet-stream")
|
|
disposition = resp.headers["Content-Disposition"]
|
|
self.assertIn("attachment", disposition)
|
|
self.assertIn("prompts_backup_", disposition)
|
|
self.assertTrue(disposition.rstrip('"').endswith(".db"))
|
|
|
|
body = await resp.read()
|
|
downloaded = self._write("downloaded.db", body)
|
|
self.assertEqual(
|
|
_texts(downloaded), ["backup me 0", "backup me 1", "backup me 2"]
|
|
)
|
|
# The server-side temp file is gone once the response is complete
|
|
self.assertEqual(len(created), 1)
|
|
self.assertFalse(os.path.exists(created[0]))
|
|
|
|
async def test_backup_of_missing_database_is_404(self):
|
|
self.api.db.model.db_path = os.path.join(self.tmpdir, "nope.db")
|
|
resp = await self.client.get("/prompt_manager/backup")
|
|
self.assertEqual(resp.status, 404)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_backup_failure_is_500(self):
|
|
self._save("x")
|
|
with patch.object(self.api.db.model, "backup_database", return_value=False):
|
|
resp = await self.client.get("/prompt_manager/backup")
|
|
self.assertEqual(resp.status, 500)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("backup", data["error"].lower())
|
|
|
|
|
|
class TestRestoreRoute(BackupRestoreTestCase):
|
|
async def test_restore_replaces_database_and_reports_backup(self):
|
|
self._save("old prompt")
|
|
old_db = self.api.db
|
|
upload = self._make_db_file(["restored a", "restored b"])
|
|
|
|
resp = await self._restore(upload)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
data = await resp.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["prompt_count"], 2)
|
|
self.assertTrue(os.path.exists(data["backup_created"]))
|
|
self.assertEqual(_texts(data["backup_created"]), ["old prompt"])
|
|
self.assertIsNot(self.api.db, old_db)
|
|
self.assertEqual(self.api.db.model.db_path, self.db_path)
|
|
|
|
texts = sorted(p["text"] for p in self.api.db.search_prompts())
|
|
self.assertEqual(texts, ["restored a", "restored b"])
|
|
|
|
async def test_restore_leaves_no_temp_files_behind(self):
|
|
upload = self._make_db_file(["restored"])
|
|
created = []
|
|
real_mkstemp = tempfile.mkstemp
|
|
|
|
def spy_mkstemp(*args, **kwargs):
|
|
result = real_mkstemp(*args, **kwargs)
|
|
created.append(result[1])
|
|
return result
|
|
|
|
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
|
|
resp = await self._restore(upload)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(len(created), 1)
|
|
for suffix in ("", "-wal", "-shm", "-journal"):
|
|
self.assertFalse(os.path.exists(created[0] + suffix))
|
|
|
|
async def test_text_file_is_rejected_and_database_untouched(self):
|
|
self._save("keep me")
|
|
upload = self._write("notes.txt", b"this is not a database")
|
|
|
|
resp = await self._restore(upload)
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("SQLite", data["error"])
|
|
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
|
self.assertFalse(
|
|
any(n.startswith("prompts.db.backup_") for n in os.listdir(self.tmpdir))
|
|
)
|
|
|
|
async def test_database_without_prompts_table_is_rejected(self):
|
|
self._save("keep me")
|
|
path = os.path.join(self.tmpdir, "other.db")
|
|
conn = sqlite3.connect(path)
|
|
conn.execute("CREATE TABLE things (id INTEGER)")
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
resp = await self._restore(path)
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("prompts", data["error"])
|
|
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
|
|
|
async def test_wrong_field_name_is_rejected(self):
|
|
upload = self._make_db_file(["restored"])
|
|
resp = await self._restore(upload, field="file")
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("database_file", data["error"])
|
|
|
|
async def test_empty_upload_is_rejected(self):
|
|
upload = self._write("empty.db", b"")
|
|
resp = await self._restore(upload)
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("empty", data["error"].lower())
|
|
|
|
async def test_upload_over_the_size_limit_is_rejected_while_streaming(self):
|
|
self._save("keep me")
|
|
self.api.restore_max_bytes = 1024
|
|
upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
|
|
|
|
resp = await self._restore(upload)
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("too large", data["error"].lower())
|
|
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
|
|
|
async def test_non_multipart_request_is_rejected(self):
|
|
resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes")
|
|
self.assertIn(resp.status, (400, 500))
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|