Vito
5772de8d27
* test: run the suite through a watchdog that names a hang (#350) tests/run_suite.py discovers the same tests as unittest discover -s tests, with the same sys.path and exit codes (5 when none ran), writes 'start: <test id>' to stderr before each test and arms faulthandler.dump_traceback_later: 120 s per test (PM_TEST_HANG_LIMIT) and 120 s after the run (PM_TEST_SHUTDOWN_LIMIT), never cancelled. A hang in a test, its teardown, its event-loop shutdown or interpreter exit prints every thread's stack and exits non-zero instead of running to the job limit. The pre-push hook now runs it too. * ci: bound every job and run the Python suite through the watchdog (#350) Every job in test.yml and code-quality.yml gets timeout-minutes = ceil(1.5 x slowest of three green main runs) + 3, min 5, with the runs and times written above jobs:. The matrix and the coverage gate run python -m tests.run_suite. ci-workflow.test.mjs checks every job in the three CI workflows has a limit and that the matrix and coverage use the runner. Job names are unchanged. * docs: document the test runner, time limits and following a push (#350) * fix(test): keep unittest's warning filters and tighten the runner's edges (#350) Review follow-ups: the runner now shows warnings as unittest.main does (DeprecationWarning and ResourceWarning were hidden), rejects non-ASCII digit limits with the one-line exit-2 error, and starts each 'start:' line on its own line at default verbosity. The docs now say that class and module fixtures count toward the test before them, that the shutdown limit covers interpreter exit, and that exit 5 on zero tests is 3.12+ behaviour applied on every version. * test: decode the runner's stderr as UTF-8 on Windows too (#350) A child Python writes a piped stderr in the ANSI codepage on Windows, so the non-ASCII limit value in the usage-error test came back as U+FFFD. The helper now sets PYTHONIOENCODING=utf-8 for the child. * test(history): give the writer's non-blocking checks a 1 s ceiling (#350) py3.13 on macos-latest failed on PR #351 with writer.close(0) taking 0.186 s against a 0.1 s ceiling while another thread held the commit. The blocked path these checks rule out waits up to 3 s (release.wait(3)), so a 1 s ceiling still tells returning at once from waiting for the release, without failing on a loaded shared runner (TESTING.md, #242). * test(restore): wait for the cancelled preview's worker before teardown (#350) py3.13 on ubuntu-latest failed on PR #351 (run 40334794505) with 'Files still open at teardown: prompts.db, prompts.db-shm' in test_cancel_while_reading_the_live_library. Cancelling the preview cannot stop its executor thread, which goes on to open the live library; the test returned without waiting, so that open could land between teardown's close_all_instances() and assert_closed(). The test now waits for the worker, as test_cancel_while_staging already does.
613 lines
25 KiB
Python
613 lines
25 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
|
|
done = threading.Event()
|
|
|
|
def slow(db):
|
|
entered.set()
|
|
release.wait(WAIT)
|
|
try:
|
|
return real(db)
|
|
finally:
|
|
done.set()
|
|
|
|
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
|
|
# The cancel cannot stop the worker thread: it still opens the live
|
|
# library. Let it finish, or teardown can find the file open.
|
|
self.assertTrue(await loop.run_in_executor(None, done.wait, WAIT))
|
|
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
|