ComfyUI/tests-unit/execution_test/test_enrich_output.py

"""Tests for enrich_output_with_assets in comfy_execution/asset_enrichment.py."""
import os
import types
import unittest
from unittest.mock import MagicMock, patch


def _make_args(enable_assets: bool):
    a = types.SimpleNamespace()
    a.enable_assets = enable_assets
    return a


def _make_register_result(ref_id="ref-id-2"):
    result = MagicMock()
    result.ref.id = ref_id
    return result


# Platform-appropriate absolute base. tempfile.gettempdir() returns C:\... on
# Windows and /tmp on POSIX, so containment via commonpath behaves naturally.
_DEFAULT_BASE = os.path.join(__import__("tempfile").gettempdir(), "asset-enrichment-test-base")


def _mocked_modules(*, enable_assets=True, register_file_in_place=None, directory=_DEFAULT_BASE):
    return {
        "comfy.cli_args": MagicMock(args=_make_args(enable_assets)),
        "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=directory)),
        "app.assets.services.ingest": MagicMock(
            register_file_in_place=register_file_in_place or MagicMock(return_value=_make_register_result()),
            DependencyMissingError=type("DependencyMissingError", (Exception,), {}),
        ),
    }


def _call(output_ui, *, enable_assets=True, file_exists=True, register_result=None, directory=_DEFAULT_BASE):
    register_mock = MagicMock(return_value=register_result or _make_register_result())
    mocked = _mocked_modules(
        enable_assets=enable_assets,
        register_file_in_place=register_mock,
        directory=directory,
    )

    # Only os.path.isfile is patched — abspath/join must run natively so the
    # containment check sees real platform paths.
    with patch.dict("sys.modules", mocked), \
         patch("os.path.isfile", return_value=file_exists):
        import importlib
        import comfy_execution.asset_enrichment as mod
        importlib.reload(mod)
        return mod.enrich_output_with_assets(output_ui)


class TestEnrichOutputWithAssets(unittest.TestCase):

    def test_disabled_returns_unchanged(self):
        output = {"images": [{"filename": "a.png", "subfolder": "", "type": "output"}]}
        result = _call(output, enable_assets=False)
        self.assertNotIn("id", result["images"][0])

    def test_non_list_value_passed_through(self):
        output = {"text": "hello"}
        result = _call(output)
        self.assertEqual(result["text"], "hello")

    def test_entry_without_filename_unchanged(self):
        output = {"latent": [{"subfolder": "", "type": "output"}]}
        result = _call(output)
        self.assertNotIn("id", result["latent"][0])

    def test_entry_without_type_unchanged(self):
        output = {"data": [{"filename": "a.png", "subfolder": ""}]}
        result = _call(output)
        self.assertNotIn("id", result["data"][0])

    def test_file_not_on_disk_unchanged(self):
        output = {"images": [{"filename": "missing.png", "subfolder": "", "type": "output"}]}
        result = _call(output, file_exists=False)
        self.assertNotIn("id", result["images"][0])

    def test_unknown_type_returns_none_directory_unchanged(self):
        output = {"images": [{"filename": "a.png", "subfolder": "", "type": "unknown"}]}
        result = _call(output, directory=None)
        self.assertNotIn("id", result["images"][0])

    def test_register_injects_only_id(self):
        reg = _make_register_result(ref_id="inline-ref")
        output = {"images": [{"filename": "new.png", "subfolder": "", "type": "output"}]}
        result = _call(output, register_result=reg)
        img = result["images"][0]
        self.assertEqual(img["id"], "inline-ref")
        # Only id is injected — no asset_hash, name, preview_url, size
        self.assertNotIn("asset_hash", img)
        self.assertNotIn("name", img)
        self.assertNotIn("preview_url", img)
        self.assertNotIn("size", img)

    def test_register_called_per_entry(self):
        register_mock = MagicMock(return_value=_make_register_result())
        mocked = _mocked_modules(register_file_in_place=register_mock)
        output = {
            "images": [
                {"filename": "a.png", "subfolder": "", "type": "output"},
                {"filename": "b.png", "subfolder": "", "type": "output"},
            ]
        }

        with patch.dict("sys.modules", mocked), \
             patch("os.path.isfile", return_value=True):
            import importlib
            import comfy_execution.asset_enrichment as mod
            importlib.reload(mod)
            mod.enrich_output_with_assets(output)

        self.assertEqual(register_mock.call_count, 2)

    def test_original_entry_not_mutated(self):
        orig = {"filename": "a.png", "subfolder": "", "type": "output"}
        output = {"images": [orig]}
        _call(output)
        self.assertNotIn("id", orig)

    def test_enrichment_error_does_not_block_sibling_entries(self):
        call_count = [0]
        good_reg = _make_register_result(ref_id="good-ref")

        def register_side_effect(abs_path, name, tags):
            call_count[0] += 1
            if call_count[0] == 1:
                raise RuntimeError("boom")
            return good_reg

        mocked = _mocked_modules(register_file_in_place=register_side_effect)

        output = {
            "images": [
                {"filename": "bad.png", "subfolder": "", "type": "output"},
                {"filename": "good.png", "subfolder": "", "type": "output"},
            ]
        }

        with patch.dict("sys.modules", mocked), \
             patch("os.path.isfile", return_value=True):
            import importlib
            import comfy_execution.asset_enrichment as mod
            importlib.reload(mod)
            result = mod.enrich_output_with_assets(output)

        imgs = result["images"]
        self.assertNotIn("id", imgs[0])
        self.assertEqual(imgs[1]["id"], "good-ref")

    def test_multiple_output_keys_all_enriched(self):
        output = {
            "images": [{"filename": "a.png", "subfolder": "", "type": "output"}],
            "videos": [{"filename": "b.mp4", "subfolder": "", "type": "output"}],
        }
        result = _call(output)
        self.assertIn("id", result["images"][0])
        self.assertIn("id", result["videos"][0])

    def test_none_entry_in_list_unchanged(self):
        output = {"images": [None, {"filename": "a.png", "subfolder": "", "type": "output"}]}
        result = _call(output)
        self.assertIsNone(result["images"][0])
        self.assertIn("id", result["images"][1])

    def test_path_traversal_subfolder_skipped(self):
        register_mock = MagicMock(return_value=_make_register_result())
        mocked = _mocked_modules(register_file_in_place=register_mock)

        output = {"images": [{"filename": "passwd", "subfolder": "../../etc", "type": "output"}]}

        # Do NOT patch os.path.abspath — real resolution is required for the containment check.
        with patch.dict("sys.modules", mocked), \
             patch("os.path.isfile", return_value=True):
            import importlib
            import comfy_execution.asset_enrichment as mod
            importlib.reload(mod)
            result = mod.enrich_output_with_assets(output)

        self.assertNotIn("id", result["images"][0])
        register_mock.assert_not_called()

    def test_absolute_filename_skipped(self):
        register_mock = MagicMock(return_value=_make_register_result())
        mocked = _mocked_modules(register_file_in_place=register_mock)

        # Absolute filename — os.path.join discards earlier components when a later one is absolute.
        absolute_filename = os.path.abspath(os.sep + "etc" + os.sep + "passwd")
        output = {"images": [{"filename": absolute_filename, "subfolder": "", "type": "output"}]}

        with patch.dict("sys.modules", mocked), \
             patch("os.path.isfile", return_value=True):
            import importlib
            import comfy_execution.asset_enrichment as mod
            importlib.reload(mod)
            result = mod.enrich_output_with_assets(output)

        self.assertNotIn("id", result["images"][0])
        register_mock.assert_not_called()


if __name__ == "__main__":
    unittest.main()