Vito
28ad3a1a1f
* refactor(database): stage restores before publishing restore_from_file is now stage_restore + publish_staged (database/ restore_stage.py). The stage lives in a .pm-restore-* folder in the library's own folder, is migrated, has damaged workflow documents repaired from an inline copy or cleared, is verified with every document read and is bound by its SHA-256. Publishing keeps a unique, standalone, 0600 safety backup (microsecond names), copies the pages through the backup API as before and then truncates the WAL, which stayed database-sized after a restore. The safety backup writer moves to database/safety_backup.py so the database layer and maintenance share it. * fix(database): read the verify streams one row at a time The provenance and review-evidence checks iterated their cursors; on a managed connection iteration prefetches 128 rows, so up to 128 large values were held at once while the docstrings promised one. Both now use fetchone. * feat(backup): review a restore, run it as an Activity job, verify backups - POST /restore/preview stages an upload in the library's folder and answers what it holds beside the live library, writing nothing live; DELETE discards it; at most two stages are held (429 too-large). - POST /restore/apply rechecks the live fingerprint (409 plan-changed) and starts "Database restore": exclusive, not cancellable, indeterminate. It waits for request leases, publishes the exact staged bytes, swaps the database, prunes safety backups and logs counts, repairs and the backup. - The legacy /restore runs the same stage and publish as an exclusive hold with its old response, so a restore and the health check, Move library and every other job refuse each other. - JobRegistry.lease: Needs review decisions and maintenance and restore previews are refused while an exclusive job runs, and a restore waits for them before it switches the file. - GET /backup holds a guarded slot, verifies a standalone snapshot before streaming it and removes the closed temp file; GET /backup/status lists the safety backups on disk and this session's verified download. * docs(testing): backup and restore contract * test(system): a restore refuses a running health check The health route test relied on a legacy restore interleaving with a running check; the restore is now exclusive and refuses it, naming the check. The diagnostics backup tests follow the backup handler into py/api/backup_restore.py. * fix(backup): drop stale restore previews off the event loop A restore drops the previews reviewed against the old library; their stage files are now removed in the executor, with the old database's connections. * test(backup): wait for the backup hold to end before reading its state CI (py3.11/3.12 on Linux, every macOS) read the registry right after the client had the last byte, before the handler left its hold and removed its temp file: the job was still running. The tests now wait on that state. * test(images): the scan dialog's Back up now leaves a finished backup hold A backup download now holds a guarded registry slot while it streams (#249 S8), so the jobs list holds that finished hold after the dialog's Back up now; the spec still proves Cancel starts no scan. * fix(backup): final review round for the restore and backup routes - I1: once the stage's copy replaced the library the restore succeeded: pruning, reopening self.db, dropping old previews and closing the old object only log, so the job finishes and the legacy /restore answers 200 instead of "the library was not changed". The generation is always bumped. - M1: a backup download holds its guarded slot only while the snapshot is written; the stream runs outside it, so a stalled download blocks nothing. - M2: a restore preview leases only its read of the live library, not the upload or the staging; a restore whose drain still times out keeps the reviewed stage under its token. - M3: the reviewed-preview store changes only on the event loop; stage folders are removed in the executor. - M4: a document whose body inflates completely while only its recorded hash or only its length disagrees is re-keyed (or merged into the document holding that text) instead of dropped. - M5: a backup download is refused only when the snapshot is not SQLite or fails integrity_check; a semantic check failure still downloads, logged, sent as X-PromptManager-Verify and recorded as not verified. - M6: a refused, failed or cancelled preview removes its stage, also when the cancel lands while a worker still stages; a timer removes expired stages without another request. - M7, N1: repair reads rows one at a time and counts distinct images. - N3: the live fingerprint never commits a transaction it did not open. * test(backup): close the old library the close-failure test leaves open CI (py3.12 Linux and the coverage gate) found the old object's executor-thread connections still open at teardown: the test patches its close_all to fail, and once the object is collected close_all_instances no longer reaches it. The test now closes it itself. * fix(backup): re-review round for restore repair, apply and backup cleanup - R1: when re-keying merges an intact body into a holder whose own body was damaged, the rewritten holder leaves the damaged list, so the clear pass no longer deletes it and unlinks the images just moved to it. - R2: a losing duplicate apply removes a stage only when the store still held it, never one another apply's job has taken. - R3: a restore that times out waiting for requests keeps its review only while it is unexpired and its slot is free; otherwise the stage goes and the job says to try again. - R4: a backup download cancelled while its snapshot is written removes the snapshot once the worker finishes.
497 lines
19 KiB
Python
497 lines
19 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 asyncio
|
|
import glob
|
|
import hashlib
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
_mock_server = MagicMock()
|
|
_mock_server.PromptServer.instance.routes = MagicMock()
|
|
sys.modules.setdefault("server", _mock_server)
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from database.models import PromptModel # noqa: E402
|
|
|
|
try:
|
|
from tests.open_handles import assert_closed # noqa: E402
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from open_handles import assert_closed # noqa: E402
|
|
|
|
import aiohttp
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import AioHTTPTestCase
|
|
|
|
from database.operations import PromptDatabase
|
|
from py.api import PromptManagerAPI, _public_path
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
def _is_upload(path):
|
|
"""Not the safety backup's temp copy (database/safety_backup.py)."""
|
|
from database.safety_backup import SAFETY_TEMP_PREFIX
|
|
|
|
return not os.path.basename(path).startswith(SAFETY_TEMP_PREFIX)
|
|
|
|
|
|
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):
|
|
app_options = {}
|
|
|
|
async def get_application(self):
|
|
self.tmpdir = tempfile.mkdtemp()
|
|
self.db_path = os.path.join(self.tmpdir, "prompts.db")
|
|
|
|
app = web.Application(**self.app_options)
|
|
routes = web.RouteTableDef()
|
|
self.api = PromptManagerAPI(db=PromptDatabase(self.db_path))
|
|
self.api.add_routes(routes)
|
|
app.router.add_routes(routes)
|
|
return app
|
|
|
|
async def tearDownAsync(self):
|
|
await super().tearDownAsync() # closes the aiohttp test client
|
|
# The automatic preview worker writes into the temp tree: stop it first.
|
|
self.api.close_auto_thumbnails()
|
|
PromptModel.close_all_instances() # restore may have swapped api.db
|
|
assert_closed(self, self.tmpdir)
|
|
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("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()
|
|
# The handler removes its temp file after the last byte is sent,
|
|
# which can be just after the client holds it: wait on that state.
|
|
while self.api._jobs.active("backup") is not None or os.path.exists(created[0]):
|
|
await asyncio.sleep(0)
|
|
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 TestRestoreHoldsTheRegistry(BackupRestoreTestCase):
|
|
"""A legacy restore and a workspace job refuse each other (#269 gap)."""
|
|
|
|
def _digest(self):
|
|
return {
|
|
suffix: hashlib.sha256(Path(self.db_path + suffix).read_bytes()).hexdigest()
|
|
for suffix in ("", "-wal")
|
|
if os.path.exists(self.db_path + suffix)
|
|
}
|
|
|
|
async def test_legacy_restore_refused_while_scan_job_runs(self):
|
|
self._save("keep me")
|
|
upload = self._make_db_file(["restored"])
|
|
gate = asyncio.Event()
|
|
|
|
async def runner(job):
|
|
await gate.wait()
|
|
|
|
scan = self.api._jobs.start(
|
|
"scan", "Scan for metadata", record_changing=True, scope={}, runner=runner
|
|
)
|
|
before = self._digest()
|
|
try:
|
|
resp = await self._restore(upload)
|
|
self.assertEqual(resp.status, 409)
|
|
self.assertEqual(
|
|
await resp.json(),
|
|
{"success": False, "error": "Scan for metadata already running"},
|
|
)
|
|
self.assertEqual(self._digest(), before)
|
|
self.assertEqual(
|
|
[p["text"] for p in self.api.db.search_prompts()], ["keep me"]
|
|
)
|
|
self.assertEqual(glob.glob(self.db_path + ".backup_*"), [])
|
|
finally:
|
|
gate.set()
|
|
await asyncio.wait_for(scan.task, 5)
|
|
|
|
async def test_scan_refused_while_legacy_restore_holds(self):
|
|
upload = self._make_db_file(["restored"])
|
|
entered = threading.Event()
|
|
release = threading.Event()
|
|
model = self.api.db.model
|
|
real_verify = model.verify_database_file
|
|
|
|
def gated(path, **kwargs):
|
|
entered.set()
|
|
release.wait(10)
|
|
return real_verify(path, **kwargs)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
with patch.object(model, "verify_database_file", gated):
|
|
restore = asyncio.ensure_future(self._restore(upload))
|
|
self.assertTrue(await loop.run_in_executor(None, entered.wait, 10))
|
|
resp = await self.client.post("/prompt_manager/jobs/scan")
|
|
body = await resp.json()
|
|
listed = (await (await self.client.get("/prompt_manager/jobs")).json())[
|
|
"jobs"
|
|
]
|
|
release.set()
|
|
restored = await restore
|
|
|
|
self.assertEqual((resp.status, body["code"]), (409, "busy"))
|
|
self.assertEqual(body["running"]["title"], "Database restore (admin page)")
|
|
self.assertEqual(
|
|
[(j["kind"], j["state"], j["cancellable"]) for j in listed],
|
|
[("restore", "running", False)],
|
|
)
|
|
self.assertEqual(restored.status, 200)
|
|
self.assertEqual(self.api._jobs.list()[0].state, "finished")
|
|
self.assertIsNone(self.api._jobs.active("restore"))
|
|
|
|
async def test_rejected_upload_ends_its_hold_failed(self):
|
|
resp = await self._restore(self._write("notes.txt", b"not a database"))
|
|
self.assertEqual(resp.status, 400)
|
|
job = self.api._jobs.list()[0]
|
|
self.assertEqual((job.kind, job.state), ("restore", "failed"))
|
|
|
|
|
|
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)
|
|
# The response names the safety backup without the server's layout
|
|
backups = glob.glob(self.db_path + ".backup_*")
|
|
self.assertEqual(len(backups), 1)
|
|
self.assertEqual(data["backup_created"], _public_path(backups[0]))
|
|
self.assertNotIn(self.tmpdir, data["backup_created"])
|
|
self.assertEqual(_texts(backups[0]), ["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("tempfile.mkstemp", spy_mkstemp):
|
|
resp = await self._restore(upload)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
# One upload, in the library's folder; it and the safety backup's
|
|
# temp copy are gone afterwards.
|
|
uploads = [p for p in created if _is_upload(p)]
|
|
self.assertEqual(len(uploads), 1)
|
|
self.assertEqual(os.path.dirname(uploads[0]), self.tmpdir)
|
|
for path in created:
|
|
for suffix in ("", "-wal", "-shm", "-journal"):
|
|
self.assertFalse(os.path.exists(path + 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_restore_cap_keeps_the_legacy_shape(self):
|
|
self.api.restore_max_bytes = 1024
|
|
upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
|
|
resp = await self._restore(upload)
|
|
data = await resp.json()
|
|
self.assertEqual(resp.status, 400)
|
|
self.assertEqual(data["code"], "too-large")
|
|
self.assertNotIn(self.tmpdir, data["error"])
|
|
|
|
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()
|
|
|
|
|
|
MIB = 2**20
|
|
|
|
|
|
class TestLargeRestores(BackupRestoreTestCase):
|
|
"""A real library's backup restores: room on disk decides, not 100 MB."""
|
|
|
|
# As ComfyUI builds its app: client_max_size is --max-upload-size (100
|
|
# MB by default). The restore streams with read_chunk, which it does not
|
|
# cap.
|
|
app_options = {"client_max_size": 100 * 2**20}
|
|
|
|
def _large_db_file(self, texts, pad_bytes):
|
|
path = self._make_db_file(texts, name="large.db")
|
|
conn = sqlite3.connect(path)
|
|
try:
|
|
conn.execute("CREATE TABLE padding(b BLOB)")
|
|
conn.execute("INSERT INTO padding VALUES (zeroblob(?))", (pad_bytes,))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
return path
|
|
|
|
async def _stream_restore(self, path, chunked=False):
|
|
with open(path, "rb") as fh:
|
|
data = fh.read()
|
|
form = aiohttp.FormData()
|
|
form.add_field(
|
|
"database_file",
|
|
data,
|
|
filename=os.path.basename(path),
|
|
content_type="application/octet-stream",
|
|
)
|
|
# chunked=False must not be passed: aiohttp 3.14 then frames the body
|
|
# as chunked while also sending a Content-Length.
|
|
options = {"chunked": True} if chunked else {}
|
|
return await self.client.post("/prompt_manager/restore", data=form, **options)
|
|
|
|
def _spy_uploads(self):
|
|
"""Patch the upload's mkstemp; the list collects each temp path made."""
|
|
created = []
|
|
real_mkstemp = tempfile.mkstemp
|
|
|
|
def spy(*args, **kwargs):
|
|
result = real_mkstemp(*args, **kwargs)
|
|
created.append(result[1])
|
|
return result
|
|
|
|
return patch("tempfile.mkstemp", spy), created
|
|
|
|
def _assert_removed(self, created):
|
|
for path in created:
|
|
for suffix in ("", "-wal", "-shm", "-journal"):
|
|
self.assertFalse(os.path.exists(path + suffix), path + suffix)
|
|
|
|
async def test_a_150_mb_backup_restores(self):
|
|
self._save("before")
|
|
upload = self._large_db_file(["restored"], 150 * MIB)
|
|
self.assertGreater(os.path.getsize(upload), 150 * MIB)
|
|
spy, created = self._spy_uploads()
|
|
with spy:
|
|
resp = await self._stream_restore(upload)
|
|
data = await resp.json()
|
|
self.assertEqual(resp.status, 200, data)
|
|
self.assertEqual(
|
|
[p["text"] for p in self.api.db.search_prompts()], ["restored"]
|
|
)
|
|
self.assertEqual(len([p for p in created if _is_upload(p)]), 1)
|
|
self._assert_removed(created)
|
|
|
|
async def test_a_short_disk_is_refused_up_front(self):
|
|
from py.api import maintenance
|
|
|
|
self._save("keep me")
|
|
upload = self._make_db_file(["restored"])
|
|
spy, created = self._spy_uploads()
|
|
with spy, patch.object(maintenance, "_free_bytes", return_value=MIB):
|
|
resp = await self._restore(upload)
|
|
data = await resp.json()
|
|
self.assertEqual(resp.status, 507)
|
|
self.assertEqual(data["code"], "insufficient-space")
|
|
self.assertTrue(
|
|
data["error"].startswith("Not enough free disk space to restore a "),
|
|
data["error"],
|
|
)
|
|
self.assertTrue(data["error"].endswith("Nothing was changed."))
|
|
self.assertNotIn(self.tmpdir, data["error"])
|
|
self.assertNotIn(tempfile.gettempdir(), data["error"])
|
|
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
|
self.assertEqual(created, []) # refused before anything was written
|
|
|
|
async def test_a_short_disk_is_refused_while_streaming_without_a_length(self):
|
|
from py.api import maintenance
|
|
|
|
self._save("keep me")
|
|
upload = self._large_db_file(["restored"], 4 * MIB)
|
|
spy, created = self._spy_uploads()
|
|
# Room for 1 MiB only, and no Content-Length to refuse up front.
|
|
with spy, patch.object(maintenance, "restore_room", return_value=MIB):
|
|
resp = await self._stream_restore(upload, chunked=True)
|
|
data = await resp.json()
|
|
self.assertEqual(resp.status, 507)
|
|
self.assertEqual(data["code"], "insufficient-space")
|
|
self.assertTrue(data["error"].endswith("Nothing was changed."))
|
|
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
|
|
self.assertEqual(len([p for p in created if _is_upload(p)]), 1)
|
|
self._assert_removed(created)
|
|
|
|
|
|
class TestRestoreRoom(unittest.TestCase):
|
|
def test_room_counts_the_upload_its_staged_copy_and_the_swap(self):
|
|
from py.api import maintenance
|
|
|
|
tmp = tempfile.mkdtemp()
|
|
try:
|
|
db_path = os.path.join(tmp, "prompts.db")
|
|
free = 10 * 2**30
|
|
with (
|
|
patch.object(maintenance, "_free_bytes", return_value=free),
|
|
patch.object(maintenance, "_device", return_value="one disk"),
|
|
):
|
|
room = maintenance.restore_room(db_path)
|
|
# Upload + staged copy + WAL growth on one disk, no database yet.
|
|
self.assertEqual(room, (free - maintenance.DISK_MARGIN_BYTES) // 3)
|
|
self.assertEqual(maintenance.size_label(int(1.6 * 2**30)), "1.6 GB")
|
|
self.assertEqual(maintenance.size_label(120 * 2**20), "120.0 MB")
|
|
finally:
|
|
os.rmdir(tmp)
|