"""Scan scope (#249 Activity-2 Task 4): new and unresolved files by default. A "Scan for metadata" job started without ``scope`` (or with ``"new"``) skips the files the Library already links with a resolved link, before their metadata is read, and counts them as Already known. ``"all"`` reads every file again, as the legacy ``/prompt_manager/scan`` route always does. """ import asyncio import json import os import sqlite3 import unittest from unittest.mock import patch try: from tests.test_api_admin import AdminAPITestCase, _sse_payloads, _write_png except ImportError: # discovered with tests/ as the top-level directory from test_api_admin import AdminAPITestCase, _sse_payloads, _write_png from database import operations # noqa: E402 from database.pair_ingestion import ingest_resolved_pair # noqa: E402 from py.config import GalleryConfig # noqa: E402 from utils.prompt_pair_resolver import resolve_prompt_pair # noqa: E402 SKIPPED_DETAIL = "Skipped · already linked" ON_DISK = object() # _link_resolved: store the file's size on disk def _graph(text): return { "1": {"class_type": "CLIPTextEncode", "inputs": {"text": text}}, "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "bad"}}, "3": { "class_type": "KSampler", "inputs": {"positive": ["1", 0], "negative": ["2", 0]}, }, } def _scope_part(scope): """The per-source scope keys this module pins (other features add counters).""" return {key: scope[key] for key in ("files", "roots", "output_folders")} class ScanScopeTestCase(AdminAPITestCase): """Helpers: link files the way the scan and the monitor do, spy on reads.""" def _link_resolved(self, path, text, size=ON_DISK): """``path`` linked with a resolved pair link (what a resolved scan writes). The row stores ``size`` as its file size: by default the file's size on disk (None when the file does not exist). """ if size is ON_DISK: size = os.path.getsize(path) if os.path.exists(path) else None metadata = {"prompt": _graph(text)} if size is not None: metadata["file_info"] = {"size": size} result = ingest_resolved_pair( self.api.db, str(path), metadata, resolve_prompt_pair(metadata) ) self.assertEqual(result["status"], "added") def _link_without_pair(self, path, text): """``path`` linked to a prompt with no pair link (fallback / FIFO link).""" prompt_id = self.api.db.save_prompt(text, None, [], None, None) self.assertTrue(self.api.db.link_image_to_prompt(prompt_id, str(path))) return prompt_id def _spy_reads(self): reads = [] original = self.api._extract_comfyui_metadata def spy(path): reads.append(os.path.basename(path)) return original(path) self.api._extract_comfyui_metadata = spy return reads async def _scan(self, body=None): kwargs = {} if body is None else {"json": body} resp = await self.client.request("POST", "/prompt_manager/jobs/scan", **kwargs) data = await resp.json() self.assertEqual(resp.status, 202, data) await asyncio.wait_for(self.api._scan_job.task, 10) detail = await self.client.request( "GET", f"/prompt_manager/jobs/{data['job']['id']}" ) return (await detail.json())["job"] class TestScanScope(ScanScopeTestCase): async def test_new_scope_skips_known_resolved_files_without_reading_them(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] known = self.output_dir / "known.png" _write_png(known, prompt_text="a red square") _write_png(self.output_dir / "new.png", prompt_text="a blue circle") self._link_resolved(known, "a red square") reads = self._spy_reads() job = await self._scan() self.assertEqual(reads, ["new.png"]) self.assertEqual(job["scope"]["files"], "new") self.assertEqual(job["counters"]["known"], 1) self.assertEqual(job["counters"]["added"], 1) self.assertEqual(job["processed"], 2) self.assertEqual( job["outcomes"]["known"], [ { "path": "known.png", "root": "output", "root_index": 0, "detail": SKIPPED_DETAIL, } ], ) self.assertIn(f"= output/known.png · {SKIPPED_DETAIL}", "\n".join(job["log"])) self.assertTrue( job["log"][0].endswith( "Scan started · 1 output folder · new and unresolved files" ) ) legacy = self.api._scan_job.last_event self.assertEqual((legacy["processed"], legacy["unchanged"]), (2, 1)) self.assertTrue(job["status_line"].startswith("Processing file 2 of 2")) async def test_whole_known_batch_skips_the_executor_reads(self): from py.api import admin GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] for name in ("a.png", "b.png"): _write_png(self.output_dir / name, prompt_text=f"prompt {name}") self._link_resolved(self.output_dir / name, f"prompt {name}") _write_png(self.output_dir / "c.png", prompt_text="prompt c") reads = self._spy_reads() with patch.object(admin, "SCAN_BATCH_SIZE", 1): job = await self._scan({"scope": "new"}) self.assertEqual(reads, ["c.png"]) self.assertEqual(job["counters"]["known"], 2) self.assertEqual(job["counters"]["added"], 1) async def test_unresolved_files_are_rescanned_in_new_scope(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] names = ("fallback.png", "conflict.png", "unlinked.png") for name in names: _write_png(self.output_dir / name, prompt_text=f"text of {name}") # Linked without a resolved pair link (the monitor's fallback / FIFO link). self._link_without_pair(self.output_dir / "fallback.png", "text of fallback") # Resolved once, then also linked to a second prompt: ambiguous. conflict = self.output_dir / "conflict.png" self._link_resolved(conflict, "text of conflict.png") self._link_without_pair(conflict, "another prompt") reads = self._spy_reads() job = await self._scan() self.assertEqual(sorted(reads), sorted(names)) self.assertNotIn(SKIPPED_DETAIL, json.dumps(job["outcomes"])) async def test_all_scope_reads_everything(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] known = self.output_dir / "known.png" _write_png(known, prompt_text="a red square") self._link_resolved(known, "a red square") reads = self._spy_reads() job = await self._scan({"scope": "all"}) self.assertEqual(reads, ["known.png"]) self.assertEqual(job["scope"]["files"], "all") self.assertEqual(job["counters"]["known"], 1) self.assertNotEqual(job["outcomes"]["known"][0]["detail"], SKIPPED_DETAIL) self.assertTrue( job["log"][0].endswith("Scan started · 1 output folder · all files") ) async def test_legacy_scan_route_reads_everything(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] known = self.output_dir / "known.png" _write_png(known, prompt_text="a red square") self._link_resolved(known, "a red square") reads = self._spy_reads() resp = await self.client.request("POST", "/prompt_manager/scan", data="{}") payloads = _sse_payloads(await resp.text()) self.assertEqual(reads, ["known.png"]) self.assertEqual(self.api._scan_job.scope["files"], "all") self.assertEqual(payloads[-1]["type"], "complete") self.assertEqual(payloads[-1]["unchanged"], 1) async def test_file_replaced_at_a_known_path_is_reread_and_needs_review(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] path = self.output_dir / "out_00012.png" _write_png(path, prompt_text="the first prompt") self._link_resolved(path, "the first prompt") before = os.path.getsize(path) # ComfyUI reused the name for a new image with another prompt. _write_png(path, prompt_text="a much longer second prompt for the same name") self.assertNotEqual(os.path.getsize(path), before) reads = self._spy_reads() job = await self._scan() self.assertEqual(reads, ["out_00012.png"]) self.assertEqual(job["counters"]["review"], 1) self.assertEqual(job["counters"]["known"], 0) self.assertIn("different prompt", job["outcomes"]["review"][0]["detail"]) async def test_scanned_files_store_their_size_so_the_next_new_scan_skips(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] _write_png(self.output_dir / "fresh.png", prompt_text="a fresh prompt") first = await self._scan() self.assertEqual(first["counters"]["added"], 1) reads = self._spy_reads() second = await self._scan() self.assertEqual(reads, []) self.assertEqual(second["outcomes"]["known"][0]["detail"], SKIPPED_DETAIL) async def test_stored_null_size_is_reread(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] path = self.output_dir / "old.png" _write_png(path, prompt_text="an old prompt") self._link_resolved(path, "an old prompt", size=None) reads = self._spy_reads() job = await self._scan() self.assertEqual(reads, ["old.png"]) # The read stores the size it was missing, so the file is known next time. self.assertEqual(job["counters"]["added"], 1) self.assertEqual( self.api.db.get_known_resolved_paths({str(path): os.path.getsize(path)}), {str(path)}, ) async def test_known_set_failure_falls_back_to_reading_every_file(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] known = self.output_dir / "known.png" _write_png(known, prompt_text="a red square") self._link_resolved(known, "a red square") reads = self._spy_reads() def broken(files): raise sqlite3.OperationalError("database is locked") self.api.db.get_known_resolved_paths = broken with self.assertLogs("prompt_manager.api", level="WARNING"): job = await self._scan() self.assertEqual(job["state"], "finished") self.assertEqual(reads, ["known.png"]) self.assertIn( "Known files could not be read · reading every file", "\n".join(job["log"]) ) async def test_empty_read_body_asks_for_the_default_scope(self): """A chunked POST with no bytes can read a body: still the default.""" from py.api.job_routes import _scan_request class Request: can_read_body = True def __init__(self, text): self._text = text async def text(self): return self._text default = ({"files": "new", "roots": None}, None) self.assertEqual(await _scan_request(Request("")), default) self.assertEqual(await _scan_request(Request(" \n")), default) self.assertEqual( await _scan_request( Request( '{"scope": "all", "roots": [2, 0, 2],' ' "paths": ["output/c", "output/a", "output/c"]}' ) ), ({"files": "all", "roots": {0: "output/a", 2: "output/c"}}, None), ) request, error = await _scan_request(Request("{")) self.assertIsNone(request) self.assertIsInstance(error, str) async def test_invalid_scope_is_refused(self): for body in ({"scope": "everything"}, {"scope": 1}, ["new"]): with self.subTest(body=body): resp = await self.client.request( "POST", "/prompt_manager/jobs/scan", json=body ) self.assertEqual(resp.status, 400) data = await resp.json() self.assertFalse(data["success"]) resp = await self.client.request( "POST", "/prompt_manager/jobs/scan", data="not json" ) self.assertEqual(resp.status, 400) self.assertIsNone(self.api._scan_job) class TestPerSourceScope(ScanScopeTestCase): """``roots``: scan only the selected effective sources (#249 S4 Task 10).""" async def asyncSetUp(self): await super().asyncSetUp() self.a = self.output_dir / "a" self.b = self.output_dir / "b" self.a.mkdir() self.b.mkdir() _write_png(self.a / "in_a.png", prompt_text="a prompt in a") _write_png(self.b / "in_b.png", prompt_text="a prompt in b") GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.a), self._root(self.b)] async def _refused(self, body): resp = await self.client.request("POST", "/prompt_manager/jobs/scan", json=body) data = await resp.json() self.assertEqual(resp.status, 400, data) self.assertFalse(data["success"]) self.assertNotIn(self.tmpdir, json.dumps(data)) self.assertIsNone(self.api._scan_job) self.assertEqual(self.api._jobs.list(), []) return data async def test_scoped_scan_only_walks_selected_root(self): walked = [] original = self.api._collect_output_media_sync def spy(output_dirs): walked.append([os.path.basename(str(d)) for d in output_dirs]) return original(output_dirs) self.api._collect_output_media_sync = spy reads = self._spy_reads() job = await self._scan({"scope": "all", "roots": [1], "paths": ["output/b"]}) self.assertEqual(walked, [["b"]]) self.assertEqual(reads, ["in_b.png"]) self.assertEqual( _scope_part(job["scope"]), {"files": "all", "roots": [1], "output_folders": 1}, ) self.assertEqual( job["outcomes"]["added"], [ { "path": "in_b.png", "root": "output/b", "root_index": 1, "detail": job["outcomes"]["added"][0]["detail"], } ], ) self.assertEqual(job["counters"]["added"], 1) self.assertEqual(self.api._scan_job.last_event["directories"], ["output/b"]) async def test_duplicate_roots_are_one_selection(self): job = await self._scan( {"roots": [1, 0, 1], "paths": ["output/b", "output/a", "output/b"]} ) self.assertEqual(job["scope"]["roots"], [0, 1]) self.assertEqual(job["scope"]["output_folders"], 2) self.assertEqual(job["counters"]["added"], 2) async def test_scope_line_counts_selected_folders(self): job = await self._scan({"roots": [0], "paths": ["output/a"]}) self.assertTrue( job["log"][0].endswith( "Scan started · 1 output folder · new and unresolved files" ), job["log"][0], ) self.assertEqual(job["scope"]["output_folders"], 1) unscoped = await self._scan({"scope": "all"}) self.assertEqual( _scope_part(unscoped["scope"]), {"files": "all", "roots": None, "output_folders": 2}, ) self.assertIn("Scan started · 2 output folders · all files", unscoped["log"][0]) async def _stale(self, body): resp = await self.client.request("POST", "/prompt_manager/jobs/scan", json=body) data = await resp.json() self.assertEqual(resp.status, 409, data) self.assertEqual(data["code"], "sources-changed") self.assertFalse(data["success"]) self.assertNotIn(self.tmpdir, json.dumps(data)) self.assertIsNone(self.api._scan_job) self.assertEqual(self.api._jobs.list(), []) return data async def test_index_past_the_sources_is_409_without_job(self): await self._stale({"roots": [2], "paths": ["output/c"]}) await self._stale({"roots": [0, 64], "paths": ["output/a", "output/x"]}) async def test_stale_index_never_scans_another_source(self): """The list changed after the page read it: source 0 was removed, so old index 1 (b) is past the end and old index 0 now names b.""" GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.b)] await self._stale({"roots": [1], "paths": ["output/b"]}) await self._stale({"roots": [0], "paths": ["output/a"]}) job = await self._scan({"roots": [0], "paths": ["output/b"]}) self.assertEqual(job["outcomes"]["added"][0]["root"], "output/b") async def test_roots_need_their_paths(self): for body in ( {"roots": [0]}, {"roots": [0], "paths": None}, {"roots": [0], "paths": "output/a"}, {"roots": [0], "paths": []}, {"roots": [0, 1], "paths": ["output/a"]}, {"roots": [0], "paths": [0]}, {"roots": [0, 0], "paths": ["output/a", "output/b"]}, ): with self.subTest(body=body): await self._refused(body) async def test_bool_and_float_roots_rejected(self): for roots in ( [True], [False], [0.0], [1.5], ["0"], [None], [-1], [0, None], [], None, 0, "0", {"0": 0}, list(range(65)), ): with self.subTest(roots=roots): paths = ["output/a"] * len(roots) if isinstance(roots, list) else [] await self._refused({"roots": roots, "paths": paths}) async def test_legacy_scan_unscoped_unchanged(self): from py.api.admin import SCAN_COUNT_KEYS resp = await self.client.request("POST", "/prompt_manager/scan", data="{}") payloads = _sse_payloads(await resp.text()) final = payloads[-1] self.assertEqual(set(final), set(SCAN_COUNT_KEYS) | {"type", "directories"}) self.assertEqual(final["directories"], ["output/a", "output/b"]) self.assertEqual(final["added"], 2) self.assertEqual( _scope_part(self.api._scan_job.scope), {"files": "all", "roots": None, "output_folders": 2}, ) async def test_legacy_scan_attaches_to_a_scoped_job_as_is(self): import threading from py.api.admin import SCAN_COUNT_KEYS gate = threading.Event() original = self.api._collect_output_media_sync def gated(output_dirs): gate.wait(5) return original(output_dirs) self.api._collect_output_media_sync = gated # The legacy route has decided to attach once it streams the job. attached = asyncio.Event() stream = self.api._stream_scan_job async def attaching(request, job): attached.set() return await stream(request, job) self.api._stream_scan_job = attaching resp = await self.client.request( "POST", "/prompt_manager/jobs/scan", json={"roots": [0], "paths": ["output/a"]}, ) self.assertEqual(resp.status, 202) started = (await resp.json())["job"] legacy = asyncio.ensure_future( self.client.request("POST", "/prompt_manager/scan", data="{}") ) try: await asyncio.wait_for(attached.wait(), 10) finally: gate.set() payloads = _sse_payloads(await (await legacy).text()) await asyncio.wait_for(self.api._scan_job.task, 10) self.assertEqual(self.api._scan_job.id, started["id"]) self.assertEqual(len(self.api._jobs.list()), 1) final = payloads[-1] self.assertEqual(set(final), set(SCAN_COUNT_KEYS) | {"type", "directories"}) self.assertEqual(final["directories"], ["output/a"]) self.assertEqual(final["added"], 1) async def test_scoped_scan_uses_the_sources_it_was_started_with(self): """Sources saved after the start never change what the job walks.""" resolved = [] original = self.api._get_all_output_dirs def counting(): resolved.append(1) return original() self.api._get_all_output_dirs = counting job = await self._scan({"roots": [1], "paths": ["output/b"]}) self.assertEqual(len(resolved), 1) # the request's own resolution only self.assertEqual(job["outcomes"]["added"][0]["root_index"], 1) class TestKnownResolvedPaths(ScanScopeTestCase): """``PromptDatabase.get_known_resolved_paths``: one streamed query.""" async def test_returns_the_given_paths_that_are_known_and_resolved(self): known, fallback = self.output_dir / "known.png", self.output_dir / "fb.png" self._link_resolved(known, "a red square", size=10) self._link_without_pair(fallback, "a fallback") asked = {str(known): 10, str(fallback): 10, str(self.output_dir / "n.png"): 10} self.assertEqual(self.api.db.get_known_resolved_paths(asked), {str(known)}) self.assertEqual(self.api.db.get_known_resolved_paths({}), set()) async def test_size_must_match_and_a_stored_null_size_is_not_known(self): same, other = self.output_dir / "same.png", self.output_dir / "other.png" unknown = self.output_dir / "unknown.png" self._link_resolved(same, "same size", size=10) self._link_resolved(other, "other size", size=10) self._link_resolved(unknown, "no stored size", size=None) asked = {str(same): 10, str(other): 11, str(unknown): 10} self.assertEqual(self.api.db.get_known_resolved_paths(asked), {str(same)}) # A file the walk could not stat is not known either. self.assertEqual(self.api.db.get_known_resolved_paths({str(same): None}), set()) async def test_row_without_file_path_matches_through_image_path(self): known = self.output_dir / "legacy.png" self._link_resolved(known, "a legacy row", size=10) with sqlite3.connect(self.api.db.model.db_path) as conn: conn.execute("UPDATE generated_images SET file_path=NULL") conn.close() found = self.api.db.get_known_resolved_paths({str(known): 10}) self.assertEqual(found, {str(known)}) async def test_stored_spelling_is_normalized_like_the_scanned_path(self): known = self.output_dir / "sub" / "known.png" self._link_resolved(known, "a red square", size=10) spelled = os.path.join(str(self.output_dir), "sub", "..", "sub", "known.png") with sqlite3.connect(self.api.db.model.db_path) as conn: conn.execute("UPDATE generated_images SET file_path=?", (spelled,)) conn.close() found = self.api.db.get_known_resolved_paths({str(known): 10}) self.assertEqual(found, {str(known)}) async def test_one_query_streamed_in_small_fetches(self): paths = [self.output_dir / f"{index}.png" for index in range(5)] for path in paths: self._link_resolved(path, f"prompt {path.name}", size=10) original = self.api.db.model.get_connection opened = [] def counting(): opened.append(1) return original() self.api.db.model.get_connection = counting with patch.object(operations, "KNOWN_PATHS_FETCH_SIZE", 2): found = self.api.db.get_known_resolved_paths({str(p): 10 for p in paths}) self.assertEqual(found, {str(p) for p in paths}) self.assertEqual(len(opened), 1) async def test_windows_path_case_matches_known_set(self): stored = self.output_dir / "Sub" / "Known.PNG" self._link_resolved(stored, "a red square", size=10) # A row written before paths were case-folded keeps its own spelling. with sqlite3.connect(self.api.db.model.db_path) as conn: conn.execute( "UPDATE generated_images SET file_path=?", (str(stored).upper(),) ) conn.close() asked = str(self.output_dir / "sub" / "known.png") with patch("os.path.normcase", lambda path: path.lower()): self.assertEqual(self.api.db.get_known_resolved_paths({asked: 10}), {asked}) if os.path.normcase("A") == "A": # a case-sensitive platform self.assertEqual(self.api.db.get_known_resolved_paths({asked: 10}), set()) class TestSkippedOutcome(unittest.TestCase): def test_skipped_file_is_already_known_with_its_own_detail(self): from py.api.admin import _scan_outcome root = os.path.abspath("out") group, row, line = _scan_outcome( os.path.join(root, "a", "b.png"), [(root, "output", 0)], "skipped" ) self.assertEqual(group, "known") self.assertEqual(row["detail"], SKIPPED_DETAIL) self.assertEqual(row["path"], "a/b.png") self.assertEqual(line, f"= output/a/b.png · {SKIPPED_DETAIL}") if __name__ == "__main__": unittest.main()