A Restore refused as busy now opens the same review again with Discard / Restore, and a restore job whose leases did not drain hands its review back (scope.review_kept) to the page that started it, which reopens it; every other path that abandons a token after Restore sends its DELETE, so held reviews can no longer lock restore out. The kept-review line says where the review is, and a full review store says what to do (I1). applyPlan resolves with its outcome and its job tracking now has failing-before proof (N5). Also: the safety-backup toggle says what the policy does (D104, M2); the restore picker has no type filter so safety backups can be chosen (M3); "Export all prompts JSON" names what /export holds (N1); small snapshots show KB (N3); the duplicated holding and save helpers are shared (N4); a dismissal test covers Cancel, x, Escape, a replacing dialog and leaving (M1). Refs #326
606 lines
24 KiB
Python
606 lines
24 KiB
Python
"""PR #335 final review fixes (#249 System S8): I1, M1 to M7, N1 and N3.
|
|
|
|
- I1: once the stage replaced the library, the restore is done: later steps
|
|
(pruning, the switch) only log, and the job and legacy route say so.
|
|
- M1: a backup download holds the registry only while its snapshot is made.
|
|
- M2: a restore preview leases only its live read; a busy drain keeps the
|
|
reviewed stage.
|
|
- M3: the reviewed-preview store changes only on the event loop.
|
|
- M4: a document whose body is intact but whose key or length is damaged is
|
|
re-keyed, not dropped.
|
|
- M5: a backup download is refused only for a corrupt snapshot; a semantic
|
|
check failure downloads, logged and reported.
|
|
- M6: a cancelled preview removes what it staged.
|
|
- M7 / N1 / N3: repair reads one row at a time; images counted once; the
|
|
fingerprint never commits a caller's transaction.
|
|
|
|
Ordering is driven by events and job state, never by timers.
|
|
"""
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import zlib
|
|
from contextlib import closing
|
|
from unittest.mock import patch
|
|
|
|
from aiohttp import web
|
|
|
|
from database import restore_stage
|
|
from database import workflow_documents as wd
|
|
from database.managed_connection import ManagedCursor
|
|
from database.models import PromptModel
|
|
from database.operations import PromptDatabase
|
|
from py.api import admin as admin_mod
|
|
from py.api import backup_restore, jobs
|
|
|
|
try:
|
|
from tests.test_restore_plans import RestoreTestCase, _texts
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_restore_plans import RestoreTestCase, _texts
|
|
|
|
WAIT = 60
|
|
|
|
|
|
class TestPostPublishSteps(RestoreTestCase):
|
|
"""I1: a failure after the copy never says the library was not changed."""
|
|
|
|
async def test_job_finishes_when_pruning_fails_after_the_copy(self):
|
|
self.save("old")
|
|
old_db = self.api.db
|
|
generation = self.api._db_generation
|
|
body = await self.preview(self.upload(["restored"]))
|
|
with (
|
|
patch.object(admin_mod, "_prune_safety_backups", side_effect=OSError("x")),
|
|
self.assertLogs(self.api.logger, "WARNING") as logs,
|
|
):
|
|
started = await self.apply(body["token"])
|
|
job = await self.finish(started["job"]["id"])
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertIn("Restored 1 prompt", job.status_line)
|
|
self.assertNotIn("not changed", "\n".join(job.log))
|
|
self.assertIsNot(self.api.db, old_db)
|
|
self.assertEqual(self.api._db_generation, generation + 1)
|
|
self.assertEqual(_texts(self.db_path), ["restored"])
|
|
self.assertIn("prune", "\n".join(logs.output).lower())
|
|
|
|
async def test_job_adopts_the_library_when_closing_the_old_one_fails(self):
|
|
self.save("old")
|
|
old_db = self.api.db
|
|
body = await self.preview(self.upload(["restored"]))
|
|
try:
|
|
with patch.object(old_db, "close_all", side_effect=RuntimeError("close")):
|
|
started = await self.apply(body["token"])
|
|
job = await self.finish(started["job"]["id"])
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertIsNot(self.api.db, old_db)
|
|
self.assertEqual(
|
|
[p["text"] for p in self.api.db.search_prompts()], ["restored"]
|
|
)
|
|
finally:
|
|
# The patched close left the old object's executor-thread
|
|
# connections open; close them before teardown removes the files
|
|
# (once the object is collected, close_all_instances misses it).
|
|
await asyncio.get_running_loop().run_in_executor(None, old_db.close_all)
|
|
|
|
async def test_job_finishes_when_opening_the_new_object_fails(self):
|
|
self.save("old")
|
|
generation = self.api._db_generation
|
|
body = await self.preview(self.upload(["restored"]))
|
|
with patch.object(
|
|
backup_restore, "PromptDatabase", side_effect=RuntimeError("open")
|
|
):
|
|
started = await self.apply(body["token"])
|
|
job = await self.finish(started["job"]["id"])
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
# The old object still points at the same, restored file.
|
|
self.assertEqual(
|
|
[p["text"] for p in self.api.db.search_prompts()], ["restored"]
|
|
)
|
|
self.assertEqual(self.api._db_generation, generation + 1)
|
|
|
|
async def test_legacy_restore_succeeds_when_pruning_fails_after_the_copy(self):
|
|
self.save("old")
|
|
old_db = self.api.db
|
|
with patch.object(admin_mod, "_prune_safety_backups", side_effect=OSError("x")):
|
|
resp = await self.client.post(
|
|
"/prompt_manager/restore", data=self.form(self.upload(["restored"]))
|
|
)
|
|
body = await resp.json()
|
|
self.assertEqual(resp.status, 200, body)
|
|
self.assertEqual(
|
|
body["message"], "Database restored successfully. Found 1 prompts."
|
|
)
|
|
self.assertIsNot(self.api.db, old_db)
|
|
|
|
|
|
class TestBackupHoldScope(RestoreTestCase):
|
|
"""M1 and M5: the hold covers the snapshot only; semantic issues warn."""
|
|
|
|
async def test_hold_ends_before_the_first_byte_is_sent(self):
|
|
self.save("keep")
|
|
seen = []
|
|
real = web.StreamResponse.prepare
|
|
|
|
async def prepare(response, request):
|
|
seen.append(self.api._jobs.active("backup"))
|
|
return await real(response, request)
|
|
|
|
with patch.object(web.StreamResponse, "prepare", prepare):
|
|
resp, data = await self.download_backup()
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertTrue(seen)
|
|
self.assertEqual(set(seen), {None}) # no hold while bytes are sent
|
|
[job] = self.api._jobs.list()
|
|
self.assertEqual((job.kind, job.state), ("backup", "finished"))
|
|
self.assertTrue(data.startswith(b"SQLite format 3\x00"))
|
|
|
|
async def test_a_corrupt_snapshot_is_refused(self):
|
|
self.save("keep")
|
|
with patch.object(
|
|
backup_restore, "snapshot_corruption", return_value="Integrity check failed"
|
|
):
|
|
resp = await self.client.get("/prompt_manager/backup")
|
|
self.assertEqual(resp.status, 500)
|
|
self.assertEqual(
|
|
await resp.json(),
|
|
{"success": False, "error": "Failed to create database backup"},
|
|
)
|
|
self.assertEqual(
|
|
[n for n in os.listdir(self.tmpdir) if n.startswith(".pm-")], []
|
|
)
|
|
[job] = self.api._jobs.list()
|
|
self.assertEqual(job.state, "failed")
|
|
|
|
async def test_a_semantic_check_failure_still_downloads_and_says_so(self):
|
|
self.save("keep")
|
|
reason = "Image pair association has a different positive prompt owner"
|
|
with (
|
|
patch.object(
|
|
PromptModel, "verify_database_file", return_value=(False, reason)
|
|
),
|
|
self.assertLogs(self.api.logger, "WARNING") as logs,
|
|
):
|
|
resp, data = await self.download_backup()
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.headers["X-PromptManager-Verify"], reason)
|
|
self.assertTrue(data.startswith(b"SQLite format 3\x00"))
|
|
self.assertIn(reason, "\n".join(logs.output))
|
|
status = await (await self.client.get("/prompt_manager/backup/status")).json()
|
|
self.assertFalse(status["last_download"]["verified"])
|
|
|
|
async def test_a_verified_download_says_ok(self):
|
|
self.save("keep")
|
|
resp, _data = await self.download_backup()
|
|
self.assertEqual(resp.headers["X-PromptManager-Verify"], "ok")
|
|
|
|
|
|
class TestPreviewLease(RestoreTestCase):
|
|
"""M2: upload and staging take no lease; a busy drain keeps the stage."""
|
|
|
|
async def test_staging_holds_no_lease(self):
|
|
leases = []
|
|
real = PromptModel.stage_restore
|
|
|
|
def stage(model, path):
|
|
leases.append(dict(self.api._jobs.leases))
|
|
return real(model, path)
|
|
|
|
with patch.object(PromptModel, "stage_restore", stage):
|
|
await self.preview(self.upload(["restored"]))
|
|
self.assertEqual(leases, [{}])
|
|
|
|
async def test_a_busy_drain_keeps_the_reviewed_stage(self):
|
|
self.save("live")
|
|
body = await self.preview(self.upload(["restored"]))
|
|
lease = self.api._jobs.lease("Stuck preview")
|
|
await lease.__aenter__()
|
|
try:
|
|
with patch.object(jobs, "LEASE_DRAIN_SECONDS", 0):
|
|
started = await self.apply(body["token"])
|
|
job = await self.finish(started["job"]["id"])
|
|
finally:
|
|
await lease.__aexit__(None, None, None)
|
|
self.assertEqual(job.state, "failed")
|
|
self.assertEqual(job.last_event["code"], "busy")
|
|
# S9 I1: the page that started it reopens the kept review (Restore again
|
|
# or Discard); the flag tells it so, and the line says where the review
|
|
# is, never "start the restore again" from a place that cannot.
|
|
self.assertIs(job.scope.get("review_kept"), True)
|
|
self.assertEqual(job.status_line, backup_restore.REVIEW_KEPT)
|
|
self.assertIn("reopened in the page that started", job.status_line)
|
|
self.assertEqual(len(self.stages()), 1)
|
|
# The same review applies once the requests are done.
|
|
job = await self.finish((await self.apply(body["token"]))["job"]["id"])
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertEqual(_texts(self.db_path), ["restored"])
|
|
|
|
async def test_a_review_that_cannot_be_kept_is_not_flagged_kept(self):
|
|
self.save("live")
|
|
body = await self.preview(self.upload(["restored"]))
|
|
lease = self.api._jobs.lease("Stuck preview")
|
|
await lease.__aenter__()
|
|
try:
|
|
with (
|
|
patch.object(jobs, "LEASE_DRAIN_SECONDS", 0),
|
|
patch.object(
|
|
backup_restore.RestorePlanStore, "give_back", lambda *_: False
|
|
),
|
|
):
|
|
started = await self.apply(body["token"])
|
|
job = await self.finish(started["job"]["id"])
|
|
finally:
|
|
await lease.__aexit__(None, None, None)
|
|
self.assertEqual(job.state, "failed")
|
|
self.assertNotIn("review_kept", job.scope)
|
|
self.assertEqual(job.status_line, backup_restore.REQUESTS_BUSY)
|
|
self.assertEqual(len(self.stages()), 0)
|
|
|
|
|
|
class TestStoreOnTheLoop(RestoreTestCase):
|
|
"""M3: entries leave the store on the loop; only file removal is threaded."""
|
|
|
|
async def test_expired_entries_leave_the_store_before_their_files_go(self):
|
|
clock = [1000.0]
|
|
store = self.api._restore_plan_store = backup_restore.RestorePlanStore(
|
|
clock=lambda: clock[0]
|
|
)
|
|
body = await self.preview(self.upload(["restored"]))
|
|
held = store.peek(body["token"])
|
|
entered, release = threading.Event(), threading.Event()
|
|
real = held.staged.discard
|
|
|
|
def slow_discard():
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
real()
|
|
|
|
held.staged.discard = slow_discard
|
|
clock[0] += 10_000
|
|
cleanup = asyncio.ensure_future(self.api._drop_expired_restores())
|
|
loop = asyncio.get_running_loop()
|
|
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
|
|
# The removal runs on a worker; the store already forgot the entry.
|
|
self.assertEqual(len(store), 0)
|
|
second = await self.preview(self.upload(["other"], "other.db"))
|
|
self.assertEqual(len(store), 1)
|
|
release.set()
|
|
await cleanup
|
|
self.assertEqual(len(self.stages()), 1)
|
|
store.peek(second["token"])
|
|
|
|
async def test_an_expired_preview_is_removed_without_another_request(self):
|
|
clock = [1000.0]
|
|
self.api._restore_plan_store = backup_restore.RestorePlanStore(
|
|
clock=lambda: clock[0], ttl=0
|
|
)
|
|
await self.preview(self.upload(["restored"]))
|
|
# The expiry timer fires without another request: wait on its state.
|
|
while self.stages() or len(self.api._restore_plans):
|
|
await asyncio.sleep(0)
|
|
self.assertEqual(self.stages(), [])
|
|
|
|
|
|
class TestCancelledPreview(RestoreTestCase):
|
|
"""M6: a cancelled preview leaves no stage behind."""
|
|
|
|
async def test_cancel_while_reading_the_live_library(self):
|
|
path = os.path.join(self.tmpdir, "uploads", "up.db")
|
|
self.upload(["restored"], "up.db")
|
|
entered, release = threading.Event(), threading.Event()
|
|
real = self.api._live_fingerprint
|
|
|
|
def slow(db):
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
return real(db)
|
|
|
|
with patch.object(self.api, "_live_fingerprint", side_effect=slow):
|
|
task = asyncio.ensure_future(
|
|
self.api._review_upload(
|
|
self.api.db, self.api._database_identity(self.api.db), path
|
|
)
|
|
)
|
|
loop = asyncio.get_running_loop()
|
|
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
|
|
self.assertEqual(len(self.stages()), 1)
|
|
task.cancel()
|
|
release.set()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
self.assertEqual(self.stages(), [])
|
|
self.assertEqual(len(self.api._restore_plans), 0)
|
|
|
|
async def test_cancel_while_staging(self):
|
|
path = os.path.join(self.tmpdir, "uploads", "up.db")
|
|
self.upload(["restored"], "up.db")
|
|
entered, release = threading.Event(), threading.Event()
|
|
real = PromptModel.stage_restore
|
|
done = threading.Event()
|
|
|
|
def slow(model, src):
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
try:
|
|
return real(model, src)
|
|
finally:
|
|
done.set()
|
|
|
|
with patch.object(PromptModel, "stage_restore", slow):
|
|
task = asyncio.ensure_future(
|
|
self.api._review_upload(
|
|
self.api.db, self.api._database_identity(self.api.db), path
|
|
)
|
|
)
|
|
loop = asyncio.get_running_loop()
|
|
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
|
|
task.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
release.set()
|
|
self.assertTrue(await loop.run_in_executor(None, done.wait, WAIT))
|
|
# The worker's finished stage is removed without another request.
|
|
while self.api._orphan_stages or self.stages():
|
|
await asyncio.sleep(0)
|
|
self.assertEqual(self.stages(), [])
|
|
|
|
|
|
class TestRekeyDocuments(RestoreTestCase):
|
|
"""M4: an intact body with a damaged key or length is re-keyed."""
|
|
|
|
def library(self, damage):
|
|
path = os.path.join(self.tmpdir, "uploads", "d.db")
|
|
self.upload(["x"], "d.db")
|
|
raw = b'{"nodes": [1, 2, 3]}'
|
|
digest = hashlib.sha256(raw).digest()
|
|
with closing(sqlite3.connect(path)) as conn:
|
|
sha, length = damage(digest, len(raw))
|
|
doc = conn.execute(
|
|
"INSERT INTO workflow_documents (sha256, byte_length, data) "
|
|
"VALUES (?,?,?)",
|
|
(sha, length, zlib.compress(raw)),
|
|
).lastrowid
|
|
prompt_id = conn.execute("SELECT id FROM prompts").fetchone()[0]
|
|
conn.execute(
|
|
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
|
"workflow_document_id) VALUES (?, '/x/a.png', 'a.png', ?)",
|
|
(prompt_id, doc),
|
|
)
|
|
conn.commit()
|
|
return path, raw.decode()
|
|
|
|
def staged_texts(self, staged):
|
|
with closing(sqlite3.connect(staged.path)) as conn:
|
|
return [
|
|
wd.read_stored_text(conn, None, row[0])
|
|
for row in conn.execute(
|
|
"SELECT workflow_document_id FROM generated_images"
|
|
)
|
|
]
|
|
|
|
def test_a_damaged_key_is_recomputed(self):
|
|
def flip(digest, length):
|
|
bad = bytearray(digest)
|
|
bad[0] ^= 1
|
|
return bytes(bad), length
|
|
|
|
path, text = self.library(flip)
|
|
staged = self.api.db.model.stage_restore(path)
|
|
try:
|
|
self.assertEqual(
|
|
staged.documents, {"repaired": 1, "cleared": 0, "images": 0}
|
|
)
|
|
self.assertEqual(self.staged_texts(staged), [text])
|
|
finally:
|
|
staged.discard()
|
|
|
|
def test_a_damaged_length_is_recomputed(self):
|
|
path, text = self.library(lambda digest, length: (digest, length + 7))
|
|
staged = self.api.db.model.stage_restore(path)
|
|
try:
|
|
self.assertEqual(staged.documents["repaired"], 1)
|
|
self.assertEqual(self.staged_texts(staged), [text])
|
|
finally:
|
|
staged.discard()
|
|
|
|
def test_a_key_that_another_document_holds_is_merged_into_it(self):
|
|
path, text = self.library(lambda digest, length: (b"\0" * 32, length))
|
|
with closing(sqlite3.connect(path)) as conn:
|
|
raw = text.encode()
|
|
conn.execute(
|
|
"INSERT INTO workflow_documents (sha256, byte_length, data) "
|
|
"VALUES (?,?,?)",
|
|
(hashlib.sha256(raw).digest(), len(raw), zlib.compress(raw)),
|
|
)
|
|
conn.commit()
|
|
staged = self.api.db.model.stage_restore(path)
|
|
try:
|
|
self.assertEqual(staged.documents["repaired"], 1)
|
|
self.assertEqual(self.staged_texts(staged), [text])
|
|
with closing(sqlite3.connect(staged.path)) as conn:
|
|
self.assertEqual(
|
|
conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone(),
|
|
(1,),
|
|
)
|
|
finally:
|
|
staged.discard()
|
|
|
|
|
|
class TestRepairReadsOneRowAtATime(RestoreTestCase):
|
|
"""M7: repair never iterates a managed cursor (128-row prefetch)."""
|
|
|
|
def test_repair_uses_fetchone(self):
|
|
upload = os.path.join(self.tmpdir, "uploads", "dmg.db")
|
|
os.makedirs(os.path.dirname(upload), exist_ok=True)
|
|
src = PromptDatabase(upload)
|
|
prompt_id = self.save("with a workflow", src)
|
|
image = src.link_image_to_prompt(
|
|
prompt_id, os.path.join(self.tmpdir, "a.png"), {"workflow": {"v": 1}}
|
|
)
|
|
conn = src.model.get_connection()
|
|
with conn:
|
|
text = wd.read_stored_text(
|
|
conn,
|
|
None,
|
|
conn.execute(
|
|
"SELECT workflow_document_id FROM generated_images WHERE id=?",
|
|
(image,),
|
|
).fetchone()[0],
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
|
"workflow_data) VALUES (?, '/x/b.png', 'b.png', ?)",
|
|
(prompt_id, text),
|
|
)
|
|
conn.execute("UPDATE workflow_documents SET data=x'00'")
|
|
try:
|
|
with patch.object(
|
|
ManagedCursor, "fetchmany", side_effect=AssertionError("prefetch")
|
|
):
|
|
result = wd.repair_damaged_documents(conn)
|
|
self.assertEqual(result["repaired"], 1)
|
|
finally:
|
|
src.close_all()
|
|
|
|
|
|
class TestFingerprintTransaction(RestoreTestCase):
|
|
"""N3: the fingerprint never commits a transaction it did not open."""
|
|
|
|
def test_a_callers_transaction_is_left_open(self):
|
|
conn = self.api.db.model.get_connection()
|
|
conn.execute("BEGIN")
|
|
conn.execute("INSERT INTO prompts (text, hash) VALUES ('pending', 'p')")
|
|
restore_stage.live_fingerprint(conn)
|
|
self.assertTrue(conn.in_transaction)
|
|
conn.rollback()
|
|
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], [])
|
|
|
|
|
|
class TestRereviewFixes(RestoreTestCase):
|
|
"""PR #335 re-review: R1 to R4."""
|
|
|
|
def test_r1_a_merge_into_a_damaged_holder_keeps_both_images(self):
|
|
rekey = TestRekeyDocuments.library
|
|
path, text = rekey(self, lambda digest, length: (b"\0" * 32, length))
|
|
raw = text.encode()
|
|
with closing(sqlite3.connect(path)) as conn:
|
|
holder = conn.execute(
|
|
"INSERT INTO workflow_documents (sha256, byte_length, data) "
|
|
"VALUES (?,?,?)",
|
|
(hashlib.sha256(raw).digest(), len(raw), b"garbage-not-zlib"),
|
|
).lastrowid
|
|
prompt_id = conn.execute("SELECT id FROM prompts").fetchone()[0]
|
|
conn.execute(
|
|
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
|
"workflow_document_id) VALUES (?, '/x/b.png', 'b.png', ?)",
|
|
(prompt_id, holder),
|
|
)
|
|
conn.commit()
|
|
staged = self.api.db.model.stage_restore(path)
|
|
try:
|
|
self.assertEqual(
|
|
staged.documents, {"repaired": 2, "cleared": 0, "images": 0}
|
|
)
|
|
self.assertEqual(
|
|
TestRekeyDocuments.staged_texts(self, staged), [text, text]
|
|
)
|
|
with closing(sqlite3.connect(staged.path)) as conn:
|
|
self.assertEqual(
|
|
conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone(),
|
|
(1,),
|
|
)
|
|
finally:
|
|
staged.discard()
|
|
|
|
async def test_r2_a_losing_duplicate_apply_leaves_the_taken_stage(self):
|
|
self.save("live")
|
|
body = await self.preview(self.upload(["restored"]))
|
|
store = self.api._restore_plans
|
|
held = store.peek(body["token"])
|
|
entered, release = threading.Event(), threading.Event()
|
|
real = self.api._live_fingerprint
|
|
|
|
def gated(db):
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
return real(db)
|
|
|
|
with patch.object(self.api, "_live_fingerprint", side_effect=gated):
|
|
second = asyncio.ensure_future(
|
|
self.client.post(
|
|
"/prompt_manager/restore/apply", json={"token": body["token"]}
|
|
)
|
|
)
|
|
loop = asyncio.get_running_loop()
|
|
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
|
|
# Meanwhile the first apply's job took the stage, and a write landed.
|
|
store.take(body["token"], "first-job")
|
|
self.save("a run count")
|
|
release.set()
|
|
resp = await second
|
|
self.assertEqual(resp.status, 409)
|
|
self.assertTrue(os.path.isdir(held.staged.directory))
|
|
await self._discard_quietly(held)
|
|
|
|
async def _discard_quietly(self, held):
|
|
await asyncio.get_running_loop().run_in_executor(None, held.staged.discard)
|
|
|
|
def test_r3_give_back_refuses_a_full_store_or_an_expired_review(self):
|
|
clock = [1000.0]
|
|
store = backup_restore.RestorePlanStore(clock=lambda: clock[0], max_stages=1)
|
|
first = store.issue(_Stage(), "f", "id", {})
|
|
held = store.take(first, "job")
|
|
other = store.issue(_Stage(), "f", "id", {})
|
|
self.assertFalse(store.give_back(first, held)) # the slot was reused
|
|
self.assertEqual(len(store), 1)
|
|
store.pop(other)
|
|
clock[0] = held.expires_at
|
|
self.assertFalse(store.give_back(first, held)) # expired meanwhile
|
|
clock[0] -= 1
|
|
self.assertTrue(store.give_back(first, held))
|
|
self.assertIs(store.peek(first), held)
|
|
|
|
async def test_r4_a_cancelled_backup_removes_its_snapshot(self):
|
|
from aiohttp.test_utils import make_mocked_request
|
|
|
|
self.save("keep")
|
|
entered, release, done = (threading.Event() for _ in range(3))
|
|
real = backup_restore.write_backup_snapshot
|
|
|
|
def slow(model):
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
try:
|
|
return real(model)
|
|
finally:
|
|
done.set()
|
|
|
|
with patch.object(backup_restore, "write_backup_snapshot", slow):
|
|
task = asyncio.ensure_future(
|
|
self.api._stream_backup(
|
|
make_mocked_request("GET", "/prompt_manager/backup")
|
|
)
|
|
)
|
|
loop = asyncio.get_running_loop()
|
|
self.assertTrue(await loop.run_in_executor(None, entered.wait, WAIT))
|
|
task.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
release.set()
|
|
self.assertTrue(await loop.run_in_executor(None, done.wait, WAIT))
|
|
|
|
async def snapshot_gone():
|
|
while any(n.startswith(".pm-safety-") for n in os.listdir(self.tmpdir)):
|
|
await asyncio.sleep(0)
|
|
|
|
# A ceiling only a regression reaches (the file never goes).
|
|
await asyncio.wait_for(snapshot_gone(), WAIT)
|
|
self.assertIsNone(self.api._jobs.active("backup"))
|
|
|
|
|
|
class _Stage:
|
|
def discard(self):
|
|
pass
|