ComfyAssets-ComfyUI_PromptManager/tests/test_history_current.py

"""Current locations require a native namespace and independent authorization."""

import importlib
import os
import sqlite3
import subprocess
import sys
import tempfile
import unittest
from contextlib import ExitStack
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

from tests.test_api_images import make_png
from database.operations import PromptDatabase
from py.config import GalleryConfig
from utils.execution_outputs import snapshot_output_roots


class CurrentLocationTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.base = Path(self.tmp.name).resolve()
        self.output = self.base / "output"
        self.temp = self.base / "temp"
        self.output.mkdir()
        self.temp.mkdir()
        self.roots = [self.output]
        self.host_output = str(self.output)
        self.host_temp = str(self.temp)
        self.stack = ExitStack()
        self.addCleanup(self.stack.close)
        self.stack.enter_context(
            patch(
                "folder_paths.get_output_directory",
                lambda: self.host_output,
                create=True,
            )
        )
        self.stack.enter_context(
            patch(
                "folder_paths.get_temp_directory", lambda: self.host_temp, create=True
            )
        )
        self.stack.enter_context(
            patch.object(
                GalleryConfig, "allowed_gallery_parents", lambda: [str(self.base)]
            )
        )
        self.stack.enter_context(
            patch.object(GalleryConfig, "path_anchors", lambda: [str(self.base)])
        )
        self.api = SimpleNamespace(_get_all_output_dirs=lambda: list(self.roots))

    def module(self, name="py.api.history_current"):
        self.assertIsNotNone(
            importlib.util.find_spec(name), f"missing current-location adapter {name}"
        )
        return importlib.import_module(name)

    def reference(self, **overrides):
        result = {
            "folder_type": "output",
            "subfolder": "",
            "filename": "image.png",
            "path_style": "windows" if os.name == "nt" else "posix",
            "root_key": snapshot_output_roots(
                "windows" if os.name == "nt" else "posix",
                self.host_output,
                self.host_temp,
            ).output_key,
        }
        result.update(overrides)
        return result

    def resolve(self, reference=None, gallery=None):
        mod = self.module()
        scope = mod.snapshot_current_output_scope(self.api)
        return mod.resolve_current_candidate(
            reference or self.reference(), scope, gallery or {}
        ).public

    def test_lexical_alias_is_not_replaced_by_canonical_root_key(self):
        mod = self.module()
        path = make_png(self.output / "image.png")
        alias = self.base / "alias"
        alias.symlink_to(self.output, target_is_directory=True)
        self.host_output = str(alias)
        reference = self.reference()
        scope = mod.snapshot_current_output_scope(self.api)
        with patch.object(
            Path, "stat", side_effect=AssertionError("locator performed IO")
        ):
            self.assertEqual(
                mod.candidate_locator(reference, scope),
                os.path.normcase(str(alias / path.name)),
            )
        candidate = mod.resolve_current_candidate(reference, scope, {}).public
        self.assertEqual(candidate["status"], "available")
        self.assertEqual(candidate["gallery_match"], "none")
        self.assertRegex(candidate["location_key"], r"^[0-9a-f]{64}$")
        wrong = self.reference(
            root_key=snapshot_output_roots(
                "windows" if os.name == "nt" else "posix", str(self.output), None
            ).output_key
        )
        self.assertEqual(
            mod.resolve_current_candidate(wrong, scope, {}).public["reason"],
            "namespace_mismatch",
        )

    def test_null_foreign_relocated_and_unsafe_descriptors_have_allowed_control(self):
        make_png(self.output / "image.png")
        self.assertEqual(self.resolve()["status"], "available")
        for changes, reason in (
            ({"root_key": None}, "root_unavailable"),
            ({"root_key": "a" * 64}, "namespace_mismatch"),
            (
                {
                    "path_style": "posix" if os.name == "nt" else "windows",
                    "subfolder": "sub",
                    "root_key": "b" * 64,
                },
                "foreign_path_style",
            ),
            ({"subfolder": "../outside"}, "namespace_mismatch"),
            ({"filename": "/absolute.png"}, "namespace_mismatch"),
        ):
            with self.subTest(changes=changes):
                candidate = self.resolve(self.reference(**changes))
                self.assertEqual(
                    (candidate["status"], candidate["reason"]), ("unavailable", reason)
                )
                self.assertNotIn("media", candidate)

    def test_temp_requires_current_gallery_authorization(self):
        make_png(self.temp / "image.png")
        reference = self.reference(
            folder_type="temp",
            root_key=snapshot_output_roots(
                "windows" if os.name == "nt" else "posix",
                self.host_output,
                self.host_temp,
            ).temp_key,
        )
        self.assertEqual(self.resolve(reference)["reason"], "not_authorized")
        self.roots.append(self.temp)
        self.assertEqual(self.resolve(reference)["status"], "available")

    def test_supported_regular_readable_control_and_distinct_failures(self):
        path = make_png(self.output / "image.png")
        self.assertEqual(self.resolve()["status"], "available")
        path.unlink()
        self.assertEqual(self.resolve()["reason"], "missing")
        path.mkdir()
        self.assertEqual(self.resolve()["reason"], "not_regular_file")
        path.rmdir()
        path.write_bytes(b"media")
        with patch("os.access", return_value=False):
            self.assertEqual(self.resolve()["reason"], "unreadable")
        (self.output / "notes.txt").write_text("private")
        self.assertEqual(
            self.resolve(self.reference(filename="notes.txt"))["reason"],
            "unsupported_media",
        )

    def test_symlink_escape_never_authorizes_another_namespace(self):
        allowed = make_png(self.output / "image.png")
        outside = make_png(self.temp / "outside.png")
        self.roots.append(self.temp)
        self.assertEqual(self.resolve()["status"], "available")
        allowed.unlink()
        allowed.symlink_to(outside)
        self.assertEqual(self.resolve()["reason"], "not_authorized")

    def test_scope_tracks_configuration_targets_and_policy_but_not_bytes(self):
        mod = self.module()
        path = make_png(self.output / "image.png")
        first = mod.snapshot_current_output_scope(self.api)
        location = mod.resolve_current_candidate(self.reference(), first, {}).public[
            "location_key"
        ]
        path.write_bytes(b"replacement")
        same = mod.snapshot_current_output_scope(self.api)
        self.assertEqual(same.current_scope, first.current_scope)
        self.assertEqual(
            mod.resolve_current_candidate(self.reference(), same, {}).public[
                "location_key"
            ],
            location,
        )
        self.roots.append(self.temp)
        changed = mod.snapshot_current_output_scope(self.api)
        self.assertNotEqual(changed.current_scope, first.current_scope)
        self.roots.reverse()
        self.assertEqual(
            mod.snapshot_current_output_scope(self.api).current_scope,
            changed.current_scope,
        )
        with patch.object(
            GalleryConfig, "allowed_gallery_parents", lambda: [str(self.output)]
        ):
            self.assertNotEqual(
                mod.snapshot_current_output_scope(self.api).current_scope,
                changed.current_scope,
            )
        alias = self.base / "alias"
        alias.symlink_to(self.output, target_is_directory=True)
        self.host_output = str(alias)
        before = mod.snapshot_current_output_scope(self.api)
        alias.unlink()
        alias.symlink_to(self.temp, target_is_directory=True)
        self.assertNotEqual(
            mod.snapshot_current_output_scope(self.api).current_scope,
            before.current_scope,
        )

    def test_gallery_projection_is_snapshot_owned_exact_and_ambiguous(self):
        mod = self.module("database.history_gallery")
        db = PromptDatabase(str(self.base / "gallery.db"))
        self.addCleanup(db.close_all)
        path = make_png(self.output / "image.png")
        locator = os.path.normcase(str(path))
        prompt = db.save_prompt("first", prompt_hash="first-hash")
        other = db.save_prompt("second", prompt_hash="second-hash")
        image = db.link_image_to_prompt(prompt, str(path))
        db.link_image_to_prompt(other, str(self.temp / path.name))
        conn = db.model.get_connection()
        with self.assertRaises(Exception):
            mod.read_history_gallery(conn, (locator,))
        with db.recorded_history_snapshot():
            projected = mod.read_history_gallery(conn, (locator,))
            self.assertTrue(conn.in_transaction)
            self.assertEqual(projected[locator]["gallery_match"], "unique")
            self.assertEqual(projected[locator]["gallery_image_id"], image)
        db.link_image_to_prompt(other, str(path))
        with db.recorded_history_snapshot():
            projected = mod.read_history_gallery(conn, (locator,))
            self.assertEqual(projected[locator]["gallery_match"], "ambiguous")
            self.assertIsNone(projected[locator]["gallery_image_id"])
            self.assertIsNone(projected[locator]["saved"])
            self.assertEqual(projected[locator]["library"]["status"], "ambiguous")
            with self.assertRaises(ValueError):
                mod.read_history_gallery(
                    conn, tuple(locator + str(i) for i in range(101))
                )

    def test_gallery_display_is_bounded_without_loading_workflow_or_prompt(self):
        mod = self.module("database.history_gallery")
        db = PromptDatabase(str(self.base / "gallery.db"))
        self.addCleanup(db.close_all)
        path = make_png(self.output / "image.png")
        locator = os.path.normcase(str(path))
        prompt = db.save_prompt(
            "secret prompt" * 10000, prompt_hash="large-hash", category="猫" * 12000
        )
        image = db.link_image_to_prompt(prompt, str(path))
        with db.model.get_connection() as conn:
            conn.execute(
                "UPDATE generated_images SET workflow_data=?,prompt_metadata=? "
                "WHERE id=?",
                ("private" * 100000, "private" * 100000, image),
            )
        statements = []
        conn.set_trace_callback(statements.append)
        self.addCleanup(conn.set_trace_callback, None)
        with db.recorded_history_snapshot():
            result = mod.read_history_gallery(conn, (locator,))
        display = result[locator]["library"]
        self.assertEqual(display["status"], "limited")
        self.assertEqual(display["tags"], [])
        self.assertFalse(
            any(
                "workflow_data" in sql or "prompt_metadata" in sql or "p.text" in sql
                for sql in statements
            )
        )
        self.assertFalse(any(sql.startswith("COMMIT") for sql in statements))

    def test_gallery_queries_stop_after_two_owners_and_two_pair_associations(self):
        mod = self.module("database.history_gallery")
        db = PromptDatabase(str(self.base / "gallery.db"))
        self.addCleanup(db.close_all)
        path = make_png(self.output / "image.png")
        unique = make_png(self.output / "unique.png")
        locator = os.path.normcase(str(path))
        unique_locator = os.path.normcase(str(unique))
        prompt = db.save_prompt("owner", prompt_hash="owner")
        image = db.link_image_to_prompt(prompt, str(unique))
        with db.model.get_connection() as conn:
            for i in range(12):
                owner = conn.execute(
                    "INSERT INTO prompts(text,hash) VALUES(?,?)",
                    (str(i), str(i)),
                ).lastrowid
                pair = conn.execute(
                    "INSERT INTO prompt_pairs"
                    "(positive_prompt_id,negative_prompt_id,title) "
                    "VALUES(?,?,?)",
                    (prompt, owner, str(i)),
                ).lastrowid
                conn.execute(
                    "INSERT INTO image_prompt_pairs(image_id,pair_id,provenance) "
                    "VALUES(?,?,'{}')",
                    (image, pair),
                )
                conn.execute(
                    "INSERT INTO generated_images"
                    "(prompt_id,image_path,filename,file_path) "
                    "VALUES(?,?,'image.png',?)",
                    (owner, str(path), locator),
                )
        fetched = {"owners": 0, "pairs": 0}

        def row_factory(cursor, values):
            names = {column[0] for column in cursor.description}
            if names == {"id", "prompt_id", "category"}:
                fetched["owners"] += 1
            elif names == {"id", "positive_prompt_id", "title", "category"}:
                fetched["pairs"] += 1
            return sqlite3.Row(cursor, values)

        conn.row_factory = row_factory
        with db.recorded_history_snapshot():
            result = mod.read_history_gallery(conn, (locator, unique_locator))
        self.assertEqual(fetched, {"owners": 3, "pairs": 2})
        self.assertEqual(result[locator]["gallery_match"], "ambiguous")
        self.assertEqual(result[unique_locator]["gallery_match"], "unique")
        self.assertEqual(result[unique_locator]["library"]["status"], "limited")
        self.assertIsNone(result[unique_locator]["library"].get("pair_id"))

    def test_host_root_failure_and_removal_do_not_search_another_root(self):
        make_png(self.output / "image.png")
        make_png(self.temp / "image.png")
        self.roots.append(self.temp)
        reference = self.reference()
        self.assertEqual(self.resolve(reference)["status"], "available")
        with patch("folder_paths.get_output_directory", side_effect=OSError("private")):
            self.assertEqual(self.resolve(reference)["reason"], "root_unavailable")
        (self.output / "image.png").unlink()
        self.output.rmdir()
        self.assertEqual(self.resolve(reference)["reason"], "root_unavailable")

    def test_unusable_lexical_namespace_changes_still_change_scope(self):
        mod = self.module()
        make_png(self.output / "image.png")
        reference = self.reference()
        self.assertEqual(self.resolve(reference)["status"], "available")
        self.host_output = "relative-one"
        first = mod.snapshot_current_output_scope(self.api)
        self.assertEqual(
            mod.resolve_current_candidate(reference, first, {}).public["reason"],
            "root_unavailable",
        )
        self.host_output = "relative-two"
        self.assertNotEqual(
            mod.snapshot_current_output_scope(self.api).current_scope,
            first.current_scope,
        )

    def test_duplicate_owner_lookup_does_not_sort_the_entire_saved_population(self):
        mod = self.module("database.history_gallery")
        db = PromptDatabase(str(self.base / "gallery.db"))
        self.addCleanup(db.close_all)
        path = os.path.normcase(str(self.output / "image.png"))
        with db.model.get_connection() as conn:
            for item in range(512):
                owner = conn.execute(
                    "INSERT INTO prompts(text) VALUES(?)", (str(item),)
                ).lastrowid
                conn.execute(
                    "INSERT INTO generated_images"
                    "(prompt_id,image_path,filename,file_path) "
                    "VALUES(?,?,'image.png',?)",
                    (owner, path, path),
                )
        steps = 0

        def bounded():
            nonlocal steps
            steps += 1
            return int(steps > 20)

        conn.set_progress_handler(bounded, 100)
        try:
            with db.recorded_history_snapshot():
                result = mod.read_history_gallery(conn, (path,))
        except sqlite3.OperationalError as error:
            self.fail(f"Two-owner lookup exceeded bounded query work: {error}")
        finally:
            conn.set_progress_handler(None, 0)
        self.assertEqual(result[path]["gallery_match"], "ambiguous")

    def test_two_same_pair_links_and_later_different_pair_withhold_pair_choice(self):
        mod = self.module("database.history_gallery")
        db = PromptDatabase(str(self.base / "gallery.db"))
        self.addCleanup(db.close_all)
        path = make_png(self.output / "image.png")
        locator = os.path.normcase(str(path))
        pair = db.save_prompt_pair("owner", "negative one")
        other_pair = db.save_prompt_pair("owner", "negative two")
        owner = db.get_prompt_pair(pair)["positive_prompt_id"]
        image = db.link_image_to_prompt(owner, str(path))
        with db.model.get_connection() as conn:
            for pair_id, provenance in (
                (pair, "{}"),
                (pair, '{"source":"second"}'),
                (other_pair, "{}"),
            ):
                conn.execute(
                    "INSERT INTO image_prompt_pairs(image_id,pair_id,provenance) "
                    "VALUES(?,?,?)",
                    (image, pair_id, provenance),
                )
        with db.recorded_history_snapshot():
            result = mod.read_history_gallery(conn, (locator,))[locator]
        self.assertEqual(result["gallery_match"], "unique")
        self.assertEqual(result["gallery_image_id"], image)
        self.assertEqual(result["library"]["status"], "limited")
        self.assertIsNone(result["library"].get("pair_id"))
        self.assertIsNone(result["saved"].get("pair"))

    def test_inspection_rejects_unsupported_suffix_before_parser(self):
        service = self.module("py.api.inspection_service")
        path = make_png(self.output / "image.png")
        observed = service.inspect_current_file(
            path, source_label="output", saved=None, library={"status": "none"}
        )
        self.assertEqual(observed["file"]["media_type"], "image")
        self.assertEqual(observed["inspection"].format, "PNG")
        unsupported = self.output / "image.txt"
        unsupported.write_bytes(path.read_bytes())
        read = service.read_inspection_file
        reads = []

        def traced(target):
            reads.append(target)
            return read(target)

        with patch.object(service, "read_inspection_file", side_effect=traced):
            with self.assertRaises(service.CurrentFileChanged):
                service.inspect_current_file(
                    unsupported,
                    source_label="output",
                    saved=None,
                    library={"status": "none"},
                )
        self.assertEqual(reads, [], "Unsupported media must not reach the parser")

    def test_preview_files_in_the_host_temp_folder_keep_preview_only_stages(self):
        # Saved raw image plus a FaceDetailer that only previews (Codex P2).
        import json

        from PIL import Image
        from PIL.PngImagePlugin import PngInfo
        from tests import inspection_parity as parity

        service = self.module("py.api.inspection_service")
        graph = parity.build("f19")["prompt"]
        graph["9"]["class_type"] = "PreviewImage"
        graph["14"] = {"class_type": "SaveImage", "inputs": {"images": ["8", 0]}}
        info = PngInfo()
        info.add_text("prompt", json.dumps(graph))
        observed = {}
        for name, folder in (("output", self.output), ("temp", self.temp)):
            path = folder / "nested" / "same.png"
            path.parent.mkdir()
            with Image.new("RGB", (4, 4)) as image:
                image.save(path, pnginfo=info)
            observed[name] = service.inspect_current_file(
                path, source_label=name, saved=None, library={"status": "none"}
            )
        saved = observed["output"]["prompts"]
        self.assertEqual(saved["positive"]["status"], "resolved")
        self.assertEqual(saved["positive"]["text"], parity.POS)
        self.assertEqual(observed["output"]["resolution"]["sampler_id"], "3")
        preview = observed["temp"]["prompts"]
        self.assertEqual(preview["positive"]["status"], "ambiguous")
        self.assertIsNone(preview["positive"]["text"])
        # An unknown host temp folder treats the file as saved output.
        self.host_temp = None
        again = service.inspect_current_file(
            self.output / "nested" / "same.png",
            source_label="output",
            saved=None,
            library={"status": "none"},
        )
        self.assertEqual(again["prompts"]["positive"]["status"], "resolved")


class HistoryImportTests(unittest.TestCase):
    def test_api_import_does_not_require_a_running_host_server(self):
        script = "; ".join(
            (
                "import sys, types",
                "sys.modules['folder_paths'] = types.ModuleType('folder_paths')",
                "import py.api",
            )
        )
        result = subprocess.run(
            [sys.executable, "-c", script],
            cwd=Path(__file__).resolve().parents[1],
            capture_output=True,
            text=True,
            check=False,
        )
        self.assertEqual(result.returncode, 0, result.stderr)

    def test_minimal_host_fixture_restores_getters_after_real_route_control(self):
        program = """
import io
from pathlib import Path
import sys
import types
import unittest

host = types.ModuleType("folder_paths")
original = lambda: "unused-minimal-host-output"
host.get_output_directory = original
sys.modules["folder_paths"] = host
from tests.test_api_history_outputs import HistoryOutputAPITests
case = HistoryOutputAPITests(
    "test_complete_zero_incomplete_unavailable_and_not_captured_are_distinct"
)
try:
    runner = unittest.TextTestRunner(stream=io.StringIO())
    result = runner.run(unittest.TestSuite([case]))
    assert host.get_output_directory is original, "fixture leaked host getter"
    assert not hasattr(host, "get_temp_directory"), "fixture leaked optional getter"
    assert result.wasSuccessful(), result.errors
finally:
    if hasattr(case, "api"):
        case.api.db.close_all()
    if hasattr(case, "_anchor_patch"):
        case._anchor_patch.stop()
    if hasattr(case, "_temp_db"):
        for suffix in ("", "-wal", "-shm"):
            Path(case._temp_db.name + suffix).unlink(missing_ok=True)
    if hasattr(case, "_output_tmp"):
        case._output_tmp.cleanup()
"""
        result = subprocess.run(
            [sys.executable, "-c", program],
            cwd=Path(__file__).resolve().parents[1],
            capture_output=True,
            text=True,
            check=False,
        )
        self.assertEqual(result.returncode, 0, result.stderr)