diff --git a/tests/browser/fast_reset.py b/tests/browser/fast_reset.py new file mode 100644 index 00000000..2e4b786b --- /dev/null +++ b/tests/browser/fast_reset.py @@ -0,0 +1,189 @@ +"""Run a browser-test fixture seeder off the live database, then publish once (#384). + +The seeders make hundreds of small commits. On a slow CI disk the live +database's commits queue behind WAL checkpoint syncs and a reset took over +30 s. ``ScratchSeeder.run`` instead: + +1. copies the live database into a scratch database (SQLite backup API) whose + connections never sync (a crash loses only the copy); +2. runs the unchanged seeder against an API proxy whose ``db`` is the scratch + database; everything else (files, thumbnails, caches, command state) is + the live API's, written exactly as before; +3. replaces the rows of every table the seed changed in the live database, in + one transaction on the live connection. Tables the seed left alone are not + touched, so state another fixture owns stays as it was. + +A commit by any other live connection while the seed ran fails the reset +instead of being overwritten (``PRAGMA data_version``). +""" + +from __future__ import annotations + +import hashlib +import shutil +import sqlite3 +import weakref +from pathlib import Path +from typing import Any, Callable, TypeVar + +T = TypeVar("T") + +# Set to 1 to run seeders on the live database as before #384 (the equivalence +# proof, tests/browser/reset_equivalence.py, compares both ways). +SEED_ON_LIVE_ENV = "PM_TEST_SEED_ON_LIVE" +SCRATCH_SCHEMA = "reset_seed" + + +def _quote(name: str) -> str: + return '"' + name.replace('"', '""') + '"' + + +def _tables(conn: sqlite3.Connection, schema: str = "main") -> list[str]: + rows = conn.execute( + f"SELECT name, sql FROM {schema}.sqlite_master WHERE type = 'table' " + "ORDER BY name" + ).fetchall() + for name, sql in rows: + if sql and sql.lstrip().upper().startswith("CREATE VIRTUAL"): + raise NotImplementedError(f"virtual table {name} cannot be copied") + return [name for name, _sql in rows] + + +def _columns(conn: sqlite3.Connection, table: str, schema: str = "main") -> list[str]: + """Writable columns (hidden and generated ones excluded), plus ``rowid`` + when the table has a rowid that no INTEGER PRIMARY KEY column carries.""" + info = conn.execute(f"PRAGMA {schema}.table_xinfo({_quote(table)})").fetchall() + columns = [row[1] for row in info if row[6] == 0] + keys = [row for row in info if row[5]] + sql = conn.execute( + f"SELECT sql FROM {schema}.sqlite_master WHERE type = 'table' AND name = ?", + (table,), + ).fetchone()[0] + without_rowid = sql is not None and "WITHOUT ROWID" in sql.upper() + integer_key = len(keys) == 1 and str(keys[0][2]).upper() == "INTEGER" + if not without_rowid and not integer_key and table != "sqlite_sequence": + columns.insert(0, "rowid") + return columns + + +def table_digests(conn: sqlite3.Connection, schema: str = "main") -> dict[str, str]: + """A digest of every table's rows (rowids included), by table name.""" + result = {} + for table in _tables(conn, schema): + selected = ", ".join( + _quote(c) if c != "rowid" else "rowid" + for c in _columns(conn, table, schema) + ) + digest = hashlib.sha256() + rows = conn.execute(f"SELECT {selected} FROM {schema}.{_quote(table)}") + for row in sorted(repr(tuple(row)) for row in rows): + digest.update(row.encode("utf-8", "surrogatepass")) + digest.update(b"\n") + result[table] = digest.hexdigest() + return result + + +class _SeedApi: + """The live API, except that ``db`` is the scratch database.""" + + def __init__(self, api: Any, db: Any) -> None: + object.__setattr__(self, "_api", api) + object.__setattr__(self, "db", db) + + def __getattr__(self, name: str) -> Any: + return getattr(self._api, name) + + def __setattr__(self, name: str, value: Any) -> None: + setattr(self._api, name, value) + + +class ScratchSeeder: + """Seed into a scratch copy and publish the changed tables in one commit.""" + + def __init__(self, api: Any, directory: Path, *, on_live: bool = False) -> None: + self._api = api + self._directory = Path(directory) + self._on_live = on_live + self._scratch = None + self._unsynced: "weakref.WeakSet[sqlite3.Connection]" = weakref.WeakSet() + + def run(self, seed: Callable[[Any], T]) -> T: + """Run ``seed(api)`` (blocking; call it in an executor thread).""" + if self._on_live: + return seed(self._api) + live = self._api.db.model.get_connection() + if live.in_transaction: + raise RuntimeError("the live connection has an open transaction") + version = live.execute("PRAGMA data_version").fetchone()[0] + scratch = self._scratch_database() + target = self._scratch_connection(scratch) + with target.using(): + live.backup(target) + before = table_digests(target) + result = seed(_SeedApi(self._api, scratch)) + target = self._scratch_connection(scratch) + if target.in_transaction: + raise RuntimeError("the seed left a scratch transaction open") + after = table_digests(target) + changed = [table for table in after if after[table] != before.get(table)] + if set(before) != set(after): + raise RuntimeError("the seed changed the database schema") + if live.execute("PRAGMA data_version").fetchone()[0] != version: + raise RuntimeError("the live database changed while the seed ran") + self._publish(live, scratch.model.db_path, changed) + return result + + def _scratch_database(self): + if self._scratch is None: + from database.operations import PromptDatabase + + self._directory.mkdir(parents=True, exist_ok=True) + self._scratch = PromptDatabase(str(self._directory / "seed.db")) + return self._scratch + + def _scratch_connection(self, scratch) -> sqlite3.Connection: + """The calling thread's scratch connection, set never to sync.""" + conn = scratch.model.get_connection() + if conn not in self._unsynced: + conn.execute("PRAGMA synchronous = OFF") + self._unsynced.add(conn) + return conn + + @staticmethod + def _publish(live: sqlite3.Connection, path: str, tables: list[str]) -> None: + if not tables: + return + foreign_keys = live.execute("PRAGMA foreign_keys").fetchone()[0] + # Rows are copied as they are; no cascade may reach a table left alone. + live.execute("PRAGMA foreign_keys = OFF") + live.execute(f"ATTACH DATABASE ? AS {SCRATCH_SCHEMA}", (path,)) + try: + live.execute("BEGIN IMMEDIATE") + try: + for table in tables: + columns = ", ".join( + _quote(c) if c != "rowid" else "rowid" + for c in _columns(live, table) + ) + name = _quote(table) + live.execute(f"DELETE FROM main.{name}") + live.execute( + f"INSERT INTO main.{name} ({columns}) " + f"SELECT {columns} FROM {SCRATCH_SCHEMA}.{name}" + ) + live.commit() + except BaseException: + live.rollback() + raise + finally: + live.execute(f"DETACH DATABASE {SCRATCH_SCHEMA}") + live.execute(f"PRAGMA foreign_keys = {int(foreign_keys)}") + + def close(self) -> None: + """Close every scratch connection, then remove the scratch files.""" + if self._on_live: + return + if self._scratch is not None: + self._scratch.close_all() + self._scratch = None + shutil.rmtree(self._directory, ignore_errors=True) diff --git a/tests/browser/reset_equivalence.py b/tests/browser/reset_equivalence.py new file mode 100644 index 00000000..b037e5c8 --- /dev/null +++ b/tests/browser/reset_equivalence.py @@ -0,0 +1,279 @@ +"""Prove the scratch-seeded resets leave the state the old resets left (#384). + +Run by hand on Linux (the browser server's platform): + + python -m tests.browser.reset_equivalence + +It starts the browser server three times: twice seeding on the live database +as before #384 (``PM_TEST_SEED_ON_LIVE=1``) and once with the scratch seeder. +Each server gets the same sequence: every reset kind of ``/__test__/reset`` +and ``/__test__/concept/reset``, with other fixtures' resets and API writes +between them. After every reset it reads every table's rows (by rowid), the +output and input file trees (size, SHA-256, mtime) and the response body. + +Values that differ between the two old runs are wall-clock values (created_at, +an unpinned concept "now", file mtimes); they are masked. Everything else must +be equal between the old runs and the new one. Exit status 0 means equal. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import sqlite3 +import subprocess +import sys +import tempfile +import time +import urllib.request +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +READY_TIMEOUT_S = 60 +CLOCK = "" + +RESETS = [ + ("/__test__/reset", {}), + ("/__test__/reset", {"extra_images": 2}), + ("/__test__/reset", {"extra_images": 108}), + ("/__test__/reset", {"include_video": True}), + ("/__test__/reset", {"prompt_inspector": True}), + ("/__test__/reset", {"native_library": True}), + ("/__test__/reset", {"native_library": True, "extra_prompts": 125}), + ("/__test__/reset", {"native_library": True, "extra_variants": 41}), + ("/__test__/reset", {"workspace_visual": True}), + ("/__test__/reset", {"full_inspector": True}), + ( + "/__test__/reset", + {"full_inspector": True, "include_video": True, "extra_images": 2}, + ), + ("/__test__/reset", {"chronological_history": True}), + ("/__test__/reset", {"chronological_history": True, "history_runs": 85}), + ("/__test__/concept/reset", {}), + ("/__test__/concept/reset", {"extra_images": 100}), + ("/__test__/concept/reset", {"seedless_images": 2, "short_seedless": True}), + ("/__test__/concept/reset", {"now": 1_780_000_000}), +] +# Writes by other fixtures and the API between resets; the resets must keep or +# replace them exactly as before. +DIRTY = [ + ("/__test__/tagging/reset", {}), + ("/__test__/system/reset", None), + ( + "/__test__/saved-views/raw", + {"page": "images", "name": "dirty", "state": "{}"}, + ), +] + + +class Server: + def __init__(self, on_live: bool) -> None: + self.tmp = tempfile.TemporaryDirectory(prefix="pm-equivalence-") + env = {**os.environ, "TMPDIR": self.tmp.name} + env.update(TEMP=self.tmp.name, TMP=self.tmp.name) + env.pop("PM_TEST_SEED_ON_LIVE", None) + if on_live: + env["PM_TEST_SEED_ON_LIVE"] = "1" + self.process = subprocess.Popen( + [ + sys.executable, + "-m", + "tests.browser.server", + "--port", + "0", + "--exit-on-stdin-eof", + ], + cwd=REPO_ROOT, + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + ) + line = self.process.stdout.readline().decode() + if not line.startswith("PM_TEST_SERVER_PORT="): + raise RuntimeError(f"no port line: {line!r}") + self.base = f"http://127.0.0.1:{int(line.split('=', 1)[1])}" + deadline = time.monotonic() + READY_TIMEOUT_S + while True: + try: + self.post_json("GET", "/prompt_manager/test", None) + break + except OSError: + if time.monotonic() > deadline or self.process.poll() is not None: + raise + time.sleep(0.1) + (self.root,) = Path(self.tmp.name).glob("pm-browser-*") + + def post_json(self, method: str, path: str, data) -> dict: + body = None if data is None else json.dumps(data).encode() + request = urllib.request.Request( + self.base + path, + data=body if method == "POST" else None, + method=method, + headers={"Content-Type": "application/json"} if body else {}, + ) + with urllib.request.urlopen(request, timeout=120) as response: # nosec B310 + text = response.read().decode() + try: + return json.loads(text) + except ValueError: + return {"text": text[:200]} + + def stop(self) -> None: + self.process.stdin.close() + self.process.wait(timeout=READY_TIMEOUT_S) + self.process.stdout.close() + self.tmp.cleanup() + + +def read_tables(path: Path) -> dict: + conn = sqlite3.connect(path.as_uri() + "?mode=ro", uri=True) + try: + tables = {} + names = [ + row[0] + for row in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name" + ) + ] + for name in names: + quoted = '"' + name.replace('"', '""') + '"' + columns = [row[1] for row in conn.execute(f"PRAGMA table_info({quoted})")] + try: + cursor = conn.execute(f"SELECT rowid, * FROM {quoted} ORDER BY rowid") + columns = ["rowid", *columns] + except sqlite3.OperationalError: + cursor = conn.execute(f"SELECT * FROM {quoted} ORDER BY 1") + rows = [ + [value.hex() if isinstance(value, bytes) else value for value in row] + for row in cursor + ] + tables[name] = {"columns": columns, "rows": rows} + return tables + finally: + conn.close() + + +def read_files(root: Path) -> dict: + files = {} + for top in ("output", "input"): + for path in sorted((root / top).rglob("*")): + if path.is_file(): + stat = path.stat() + files[path.relative_to(root).as_posix()] = [ + stat.st_size, + hashlib.sha256(path.read_bytes()).hexdigest(), + stat.st_mtime_ns, + ] + return files + + +def flatten(value, prefix="") -> dict: + if isinstance(value, dict): + out = {} + for key, item in value.items(): + out.update(flatten(item, f"{prefix}/{key}")) + return out + if isinstance(value, list): + out = {f"{prefix}#len": len(value)} + for index, item in enumerate(value): + out.update(flatten(item, f"{prefix}[{index}]")) + return out + return {prefix: value} + + +def snapshot(server: Server, response: dict) -> dict: + flat = {} + for table, data in read_tables(server.root / "browser.db").items(): + flat[f"table {table}#rows"] = len(data["rows"]) + for index, row in enumerate(data["rows"]): + for column, value in zip(data["columns"], row): + flat[f"table {table}[{index}].{column}"] = value + for name, (size, digest, mtime) in read_files(server.root).items(): + flat[f"file {name}"] = [size, digest] + flat[f"file {name}#mtime"] = mtime + flat["files#count"] = sum( + 1 for key in flat if key.startswith("file ") and "#" not in key + ) + for key, value in flatten(response).items(): + flat[f"response{key}"] = value + # Each server has its own temp root and a random storage epoch; name them + # instead of masking every column that carries them. + epoch = flat.get("table execution_history_metadata[0].storage_epoch") + names = [(str(server.root), "")] + if isinstance(epoch, str) and epoch: + names.append((epoch, "")) + for key, value in flat.items(): + if isinstance(value, str): + for text, name in names: + value = value.replace(text, name) + flat[key] = value + return flat + + +def run(on_live: bool) -> list[dict]: + server = Server(on_live) + try: + states = [] + for index, (path, data) in enumerate(RESETS * 2): + dirty_path, dirty_data = DIRTY[index % len(DIRTY)] + server.post_json("POST", dirty_path, dirty_data) + started = time.monotonic() + response = server.post_json("POST", path, data) + elapsed = time.monotonic() - started + print(f" {'old' if on_live else 'new'} {path} {data}: {elapsed:.2f}s") + states.append(snapshot(server, response)) + return states + finally: + server.stop() + + +def column_key(key: str) -> str: + """``table t[3].created_at`` -> ``table t.created_at`` (mask whole columns).""" + if key.startswith("table ") and "[" in key: + head, _, tail = key.partition("[") + return head + tail.partition("]")[2] + return key + + +def main() -> int: + old_a, old_b, new = run(True), run(True), run(False) + failures = 0 + masked_all = set() + for step, (path, data) in enumerate(RESETS * 2): + a, b, n = old_a[step], old_b[step], new[step] + masked = { + column_key(key) + for key in set(a) | set(b) + if a.get(key, CLOCK) != b.get(key, CLOCK) + } + masked_all |= masked + keys = set(a) | set(n) + diffs = [ + key + for key in sorted(keys) + if column_key(key) not in masked + and a.get(key, "") != n.get(key, "") + ] + tables = len( + {k.split("[")[0].split("#")[0] for k in a if k.startswith("table ")} + ) + status = "EQUAL" if not diffs else f"DIFFERENT ({len(diffs)} values)" + print( + f"{status}: {path} {json.dumps(data)}: {tables} tables, " + f"{a['files#count']} files, {len(keys)} values, " + f"{len(masked)} masked columns" + ) + for key in diffs[:20]: + old, now = a.get(key, ""), n.get(key, "") + print(f" {key}: old {old!r} new {now!r}") + failures += bool(diffs) + print("masked (differ between the two old runs):") + for key in sorted(masked_all): + print(f" {key}") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/browser/server.py b/tests/browser/server.py index 707634de..3c31b707 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -21,6 +21,7 @@ from PIL.PngImagePlugin import PngInfo from tests.browser.request_hold import register_request_holds from tests.browser import slow_reset_watch +from tests.browser.fast_reset import SEED_ON_LIVE_ENV, ScratchSeeder from tests.browser.workspace_activity_fixture import register_activity_commands from tests.browser.workspace_system_fixture import register_system_commands from tests.browser.workspace_integrations_fixture import ( @@ -225,6 +226,10 @@ def create_app(root: Path) -> web.Application: register_integrations_commands(app, routes, api, root) request_holds = register_request_holds(app, routes) slow_reset_watch.install(app, api) + # Resets seed a scratch copy and publish it in one commit (#384). + seeder = ScratchSeeder( + api, root / "reset-scratch", on_live=os.environ.get(SEED_ON_LIVE_ENV) == "1" + ) async def canonical_asset(request: web.Request) -> web.Response: filename, content_type = CANONICAL_ASSETS[request.path] @@ -238,6 +243,7 @@ def create_app(root: Path) -> web.Application: routes.get(path)(canonical_asset) def seed( + target, extra_images: int = 0, include_video: bool = False, prompt_inspector: bool = False, @@ -245,7 +251,7 @@ def create_app(root: Path) -> web.Application: extra_prompts: int = 0, extra_variants: int = 0, ) -> None: - with api.db.model.get_connection() as conn: + with target.db.model.get_connection() as conn: for table in ( "saved_views", "image_favorites", @@ -308,28 +314,30 @@ def create_app(root: Path) -> web.Application: image.save(image_path, pnginfo=info) stamp = time.time() - index * 10 os.utime(image_path, (stamp, stamp)) - pid = api.db.save_prompt( + pid = target.db.save_prompt( text, tags=[tag], prompt_hash=generate_prompt_hash(text) ) - api.db.link_image_to_prompt(pid, str(image_path)) + target.db.link_image_to_prompt(pid, str(image_path)) if index == 0: video_prompt_id = pid for number in range(extra_images): extra = root / directory / "bulk" / f"{number:04}.png" extra.parent.mkdir(exist_ok=True) shutil.copyfile(image_path, extra) - api.db.link_image_to_prompt(pid, str(extra)) + target.db.link_image_to_prompt(pid, str(extra)) if native_library: - positive = api.db.get_prompt_by_hash(generate_prompt_hash("Cat portrait")) + positive = target.db.get_prompt_by_hash( + generate_prompt_hash("Cat portrait") + ) positive_id = positive["id"] - api.db.update_prompt_metadata( + target.db.update_prompt_metadata( positive_id, category="Animals", tags=["cats", "studio"] ) variants = [] for index, negative in enumerate( ["blurry, watermark", "low quality", "low contrast", ""] ): - pair_id = api.db.save_prompt_pair( + pair_id = target.db.save_prompt_pair( "Cat portrait", negative, title=f"Portrait variant {index + 1}", @@ -337,28 +345,30 @@ def create_app(root: Path) -> web.Application: tags=["studio"] if index == 0 else ["outdoors"], ) variants.append(pair_id) - with api.db.model.get_connection() as conn: + with target.db.model.get_connection() as conn: images = conn.execute( "SELECT id FROM generated_images WHERE prompt_id=? ORDER BY id", (positive_id,), ).fetchall() for index, image in enumerate(images[:2]): - api.db.link_image_to_pair( + target.db.link_image_to_pair( variants[index], image["id"], {"source": "fixture", "anchor": index} ) - api.db.link_image_to_pair( + target.db.link_image_to_pair( variants[0], images[0]["id"], {"source": "fixture second evidence"} ) - api.db.record_prompt_pair_use(variants[0], times=3) + target.db.record_prompt_pair_use(variants[0], times=3) for number in range(extra_variants): - api.db.save_prompt_pair("Cat portrait", f"Negative extra {number:04}") + target.db.save_prompt_pair( + "Cat portrait", f"Negative extra {number:04}" + ) legacy = root / "output" / "legacy-unassigned.png" shutil.copyfile(root / "output" / "shared" / "same.png", legacy) - api.db.link_image_to_prompt(positive_id, str(legacy)) - api.db.save_prompt("Dog landscape", category="Animals", tags=["dogs"]) - api.db.save_prompt_pair("Dual role portrait", "Cat portrait") + target.db.link_image_to_prompt(positive_id, str(legacy)) + target.db.save_prompt("Dog landscape", category="Animals", tags=["dogs"]) + target.db.save_prompt_pair("Dual role portrait", "Cat portrait") for number in range(extra_prompts): - api.db.save_prompt( + target.db.save_prompt( f"Library extra {number:04}", category="Bulk", tags=["bulk"] ) if prompt_inspector: @@ -379,8 +389,8 @@ def create_app(root: Path) -> web.Application: from py.api.thumbnail_settings import ThumbnailSettings # 150 px previews in the default format (#317 thumbnail settings). - api._generate_gallery_thumbnails_sync( - api._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150) + target._generate_gallery_thumbnails_sync( + target._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150) ) if include_video: # Seed after thumbnail generation to exercise a real unthumbnailed video. @@ -388,8 +398,8 @@ def create_app(root: Path) -> web.Application: shutil.copyfile( Path(__file__).resolve().parents[1] / "fixtures" / "tiny.webm", video ) - api.db.link_image_to_prompt(video_prompt_id, str(video)) - api.invalidate_gallery_cache() + target.db.link_image_to_prompt(video_prompt_id, str(video)) + target.invalidate_gallery_cache() history_fixture = None @@ -434,7 +444,10 @@ def create_app(root: Path) -> web.Application: raise web.HTTPBadRequest() if chronological_history: manifest = await api._run_in_executor( - lambda: seed_workspace_history(api, root, history_runs=history_runs) + seeder.run, + lambda target: seed_workspace_history( + target, root, history_runs=history_runs + ), ) history_fixture = manifest return web.json_response({"success": True, **manifest}) @@ -492,25 +505,31 @@ def create_app(root: Path) -> web.Application: raise web.HTTPBadRequest() if workspace_visual: history_fixture = None - await api._run_in_executor(seed_workspace_visual, api, root) + await api._run_in_executor( + seeder.run, lambda target: seed_workspace_visual(target, root) + ) return web.json_response({"success": True}) if full_inspector: history_fixture = None await api._run_in_executor( - lambda: seed_workspace_inspector( - api, root, extra_images=extra, include_video=include_video - ) + seeder.run, + lambda target: seed_workspace_inspector( + target, root, extra_images=extra, include_video=include_video + ), ) return web.json_response({"success": True}) history_fixture = None await api._run_in_executor( - seed, - extra, - include_video, - prompt_inspector, - native_library, - extra_prompts, - extra_variants, + seeder.run, + lambda target: seed( + target, + extra, + include_video, + prompt_inspector, + native_library, + extra_prompts, + extra_variants, + ), ) return web.json_response({"success": True}) @@ -550,14 +569,15 @@ def create_app(root: Path) -> web.Application: reveal_stub.reset() history_fixture = None manifest = await api._run_in_executor( - lambda: seed_workspace_concept( - api, + seeder.run, + lambda target: seed_workspace_concept( + target, root, extra_images=data.get("extra_images", 0), seedless_images=data.get("seedless_images", 0), short_seedless=data.get("short_seedless", False), now=data.get("now"), - ) + ), ) return web.json_response({"success": True, **manifest}) @@ -625,9 +645,10 @@ def create_app(root: Path) -> web.Application: async def cleanup(application: web.Application) -> None: # A scan of an ambiguous fixture file captures it for Needs review. api.close_annotation_captures() + seeder.close() api.db.close_all() - seed() + seed(api) # Open Needs review change streams end first, as on the real host. app.on_shutdown.append(api.close_review_events) app.on_cleanup.append(cleanup) diff --git a/tests/test_browser_fast_reset.py b/tests/test_browser_fast_reset.py new file mode 100644 index 00000000..27cc316a --- /dev/null +++ b/tests/test_browser_fast_reset.py @@ -0,0 +1,194 @@ +"""Browser-test resets seed a scratch copy and publish it in one commit (#384).""" + +import hashlib +import sqlite3 +import tempfile +import types +import unittest +from pathlib import Path + +try: + from tests.open_handles import assert_closed +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + +from database.operations import PromptDatabase +from tests.browser.fast_reset import ScratchSeeder, table_digests + + +def digests(db: PromptDatabase) -> dict: + return table_digests(db.model.get_connection()) + + +def add_prompt(db: PromptDatabase, text: str) -> int: + with db.model.get_connection() as conn: + return conn.execute( + "INSERT INTO prompts (text, hash, created_at, updated_at) " + "VALUES (?, ?, '2026-01-01 00:00:00', '2026-01-01 00:00:00')", + (text, hashlib.sha256(text.encode()).hexdigest()), + ).lastrowid + + +def add_image(db: PromptDatabase, prompt_id: int, name: str) -> int: + with db.model.get_connection() as conn: + return conn.execute( + "INSERT INTO generated_images (prompt_id, image_path, filename, " + "generation_time) VALUES (?, ?, ?, '2026-01-01 00:00:00')", + (prompt_id, f"/out/{name}", name), + ).lastrowid + + +def reseed(api) -> str: + """A seeder in the old style: clears its tables, then many small commits.""" + db = api.db + with db.model.get_connection() as conn: + for table in ("generated_images", "prompt_tags", "tags", "prompts"): + conn.execute(f"DELETE FROM {table}") + for index in range(30): + prompt_id = add_prompt(db, f"seeded {index}") + add_image(db, prompt_id, f"seeded-{index}.png") + with db.model.get_connection() as conn: + tag_id = conn.execute( + "INSERT INTO tags (name) VALUES (?)", (f"tag{index}",) + ).lastrowid + conn.execute( + "INSERT INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)", + (prompt_id, tag_id), + ) + return "manifest" + + +class TestScratchSeeder(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.root = Path(self.tmp.name) + self.addCleanup(assert_closed, self, self.tmp.name) + self.live = self.database("live.db") + self.api = types.SimpleNamespace(db=self.live, marker="live api") + self.seeder = ScratchSeeder(self.api, self.root / "scratch") + self.addCleanup(self.seeder.close) + # Rows from "earlier tests": AUTOINCREMENT positions have moved on. + for index in range(5): + add_image(self.live, add_prompt(self.live, f"old {index}"), f"o{index}") + with self.live.model.get_connection() as conn: + conn.execute("DELETE FROM prompts WHERE text = 'old 4'") + + def database(self, name: str) -> PromptDatabase: + db = PromptDatabase(str(self.root / name)) + self.addCleanup(db.close_all) + return db + + def oracle(self, seed): + """Run ``seed`` the old way, on a copy of the live database.""" + path = self.root / "oracle.db" + target = sqlite3.connect(path) + try: + self.live.model.get_connection().backup(target) + finally: + target.close() + db = self.database("oracle.db") + seed(types.SimpleNamespace(db=db)) + return digests(db) + + def test_publishes_exactly_what_the_seed_writes_on_live(self): + expected = self.oracle(reseed) + self.assertEqual(self.seeder.run(reseed), "manifest") + self.assertEqual(digests(self.live), expected) + with self.live.model.get_connection() as conn: + first = conn.execute("SELECT MIN(id) FROM prompts").fetchone()[0] + self.assertEqual(first, 6, "AUTOINCREMENT continues as on live") + + def test_runs_twice_from_the_state_the_first_run_left(self): + self.seeder.run(reseed) + expected = self.oracle(reseed) + self.seeder.run(reseed) + self.assertEqual(digests(self.live), expected) + + def test_tables_the_seed_did_not_change_keep_their_rows(self): + before = digests(self.live) + + def rename(api): + with api.db.model.get_connection() as conn: + conn.execute("UPDATE prompts SET text = 'renamed' WHERE id = 1") + + self.seeder.run(rename) + after = digests(self.live) + changed = {table for table in after if after[table] != before[table]} + self.assertEqual(changed, {"prompts"}) + with self.live.model.get_connection() as conn: + images = conn.execute("SELECT COUNT(*) FROM generated_images").fetchone() + # Replacing prompts' rows never cascades into generated_images. + self.assertEqual(images[0], 4) + + def test_the_seed_sees_the_scratch_database_and_the_live_api(self): + seen = {} + + def seed(api): + seen["db"] = api.db + seen["marker"] = api.marker + api.note = "set by the seed" + + self.seeder.run(seed) + self.assertIsNot(seen["db"], self.live) + self.assertNotEqual(seen["db"].model.db_path, self.live.model.db_path) + self.assertEqual(seen["marker"], "live api") + self.assertEqual(self.api.note, "set by the seed") + self.assertIs(self.api.db, self.live) + + def test_a_failing_seed_leaves_live_untouched(self): + before = digests(self.live) + + def broken(api): + reseed(api) + raise ValueError("seed failed") + + with self.assertRaisesRegex(ValueError, "seed failed"): + self.seeder.run(broken) + self.assertEqual(digests(self.live), before) + self.assertIs(self.api.db, self.live) + + def test_a_concurrent_live_commit_fails_the_reset(self): + other = self.database("live.db") + + def racing(api): + reseed(api) + add_prompt(other, "written to live meanwhile") + + with self.assertRaisesRegex(RuntimeError, "changed while"): + self.seeder.run(racing) + with self.live.model.get_connection() as conn: + texts = [row[0] for row in conn.execute("SELECT text FROM prompts")] + self.assertIn("written to live meanwhile", texts) + self.assertNotIn("seeded 0", texts) + + def test_the_live_connection_is_left_as_it_was(self): + self.seeder.run(reseed) + conn = self.live.model.get_connection() + self.assertFalse(conn.in_transaction) + self.assertEqual(conn.execute("PRAGMA foreign_keys").fetchone()[0], 1) + names = [row[1] for row in conn.execute("PRAGMA database_list")] + self.assertEqual(names, ["main"]) + self.assertEqual(conn.execute("PRAGMA foreign_key_check").fetchall(), []) + + def test_on_live_mode_runs_the_seed_on_the_live_api(self): + seeder = ScratchSeeder(self.api, self.root / "unused", on_live=True) + self.addCleanup(seeder.close) + expected = self.oracle(reseed) + seen = [] + self.assertEqual( + seeder.run(lambda api: seen.append(api) or reseed(api)), "manifest" + ) + self.assertIs(seen[0], self.api) + self.assertEqual(digests(self.live), expected) + self.assertFalse((self.root / "unused").exists()) + + def test_close_releases_the_scratch_files(self): + self.seeder.run(reseed) + self.seeder.close() + assert_closed(self, self.root / "scratch") + self.seeder.close() # idempotent + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_browser_slow_reset_watch.py b/tests/test_browser_slow_reset_watch.py index 631e7ca4..569983b8 100644 --- a/tests/test_browser_slow_reset_watch.py +++ b/tests/test_browser_slow_reset_watch.py @@ -221,12 +221,12 @@ class TestSlowResetOnTheServer(unittest.TestCase): self.assertRegex( text, r"kernel: asyncio_\d+ \(0x[0-9a-f]{16}\) state \w, waiting in " ) - # The plain reset's seed() waits on its first DELETE. - self.assertRegex(text, r"line \d+ in seed\n") + # The plain reset seeds a scratch copy (#384), then waits to publish it. + self.assertRegex(text, r"line \d+ in _publish\n") self.assertIn("POST /__test__/reset answered 200 after", text) self.assertRegex( text, - r"executor call create_app..seed: queued \d+\.\d+s, ran \d+\.\d+s", + r"executor call ScratchSeeder.run: queued \d+\.\d+s, ran \d+\.\d+s", )