mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-25 16:37:23 +08:00
1ebbee4522
Some checks are pending
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Build package / Build Test (3.10) (push) Waiting to run
Build package / Build Test (3.11) (push) Waiting to run
Build package / Build Test (3.12) (push) Waiting to run
Build package / Build Test (3.13) (push) Waiting to run
Build package / Build Test (3.14) (push) Waiting to run
The previous test setup patched os.path.abspath to identity and used a POSIX-style '/output' base, which collided with Windows path separators in os.path.commonpath. Drop the abspath/join patches and use a real tempdir-rooted base so the containment check runs against actual platform paths.
246 lines
10 KiB
Python
246 lines
10 KiB
Python
"""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_db_ref(ref_id="ref-id-1"):
|
|
ref = MagicMock()
|
|
ref.id = ref_id
|
|
return ref
|
|
|
|
|
|
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 _call(output_ui, *, enable_assets=True, file_exists=True, db_ref=None, register_result=None, directory=_DEFAULT_BASE):
|
|
fake_session_cm = MagicMock()
|
|
fake_session_cm.__enter__ = MagicMock(return_value=MagicMock())
|
|
fake_session_cm.__exit__ = MagicMock(return_value=False)
|
|
|
|
mocked_modules = {
|
|
"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=MagicMock(return_value=register_result or _make_register_result()),
|
|
DependencyMissingError=type("DependencyMissingError", (Exception,), {}),
|
|
),
|
|
"app.assets.database.queries.asset_reference": MagicMock(
|
|
get_reference_by_file_path=MagicMock(return_value=db_ref),
|
|
),
|
|
"app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)),
|
|
}
|
|
|
|
# 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_modules), \
|
|
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_db_hit_injects_id(self):
|
|
db_ref = _make_db_ref(ref_id="db-ref")
|
|
output = {"images": [{"filename": "a.png", "subfolder": "", "type": "output"}]}
|
|
result = _call(output, db_ref=db_ref)
|
|
img = result["images"][0]
|
|
self.assertEqual(img["id"], "db-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_db_miss_falls_back_to_register(self):
|
|
reg = _make_register_result(ref_id="inline-ref")
|
|
output = {"images": [{"filename": "new.png", "subfolder": "", "type": "output"}]}
|
|
result = _call(output, db_ref=None, register_result=reg)
|
|
img = result["images"][0]
|
|
self.assertEqual(img["id"], "inline-ref")
|
|
self.assertNotIn("asset_hash", img)
|
|
self.assertNotIn("name", img)
|
|
|
|
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
|
|
|
|
fake_session_cm = MagicMock()
|
|
fake_session_cm.__enter__ = MagicMock(return_value=MagicMock())
|
|
fake_session_cm.__exit__ = MagicMock(return_value=False)
|
|
|
|
mocked_modules = {
|
|
"comfy.cli_args": MagicMock(args=_make_args(True)),
|
|
"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)),
|
|
"app.assets.services.ingest": MagicMock(
|
|
register_file_in_place=register_side_effect,
|
|
DependencyMissingError=type("DependencyMissingError", (Exception,), {}),
|
|
),
|
|
"app.assets.database.queries.asset_reference": MagicMock(
|
|
get_reference_by_file_path=MagicMock(return_value=None),
|
|
),
|
|
"app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)),
|
|
}
|
|
|
|
output = {
|
|
"images": [
|
|
{"filename": "bad.png", "subfolder": "", "type": "output"},
|
|
{"filename": "good.png", "subfolder": "", "type": "output"},
|
|
]
|
|
}
|
|
|
|
with patch.dict("sys.modules", mocked_modules), \
|
|
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):
|
|
fake_session_cm = MagicMock()
|
|
fake_session_cm.__enter__ = MagicMock(return_value=MagicMock())
|
|
fake_session_cm.__exit__ = MagicMock(return_value=False)
|
|
|
|
register_mock = MagicMock(return_value=_make_register_result())
|
|
mocked_modules = {
|
|
"comfy.cli_args": MagicMock(args=_make_args(True)),
|
|
"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)),
|
|
"app.assets.services.ingest": MagicMock(
|
|
register_file_in_place=register_mock,
|
|
DependencyMissingError=type("DependencyMissingError", (Exception,), {}),
|
|
),
|
|
"app.assets.database.queries.asset_reference": MagicMock(
|
|
get_reference_by_file_path=MagicMock(return_value=None),
|
|
),
|
|
"app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)),
|
|
}
|
|
|
|
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_modules), \
|
|
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):
|
|
fake_session_cm = MagicMock()
|
|
fake_session_cm.__enter__ = MagicMock(return_value=MagicMock())
|
|
fake_session_cm.__exit__ = MagicMock(return_value=False)
|
|
|
|
register_mock = MagicMock(return_value=_make_register_result())
|
|
mocked_modules = {
|
|
"comfy.cli_args": MagicMock(args=_make_args(True)),
|
|
"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)),
|
|
"app.assets.services.ingest": MagicMock(
|
|
register_file_in_place=register_mock,
|
|
DependencyMissingError=type("DependencyMissingError", (Exception,), {}),
|
|
),
|
|
"app.assets.database.queries.asset_reference": MagicMock(
|
|
get_reference_by_file_path=MagicMock(return_value=None),
|
|
),
|
|
"app.database.db": MagicMock(create_session=MagicMock(return_value=fake_session_cm)),
|
|
}
|
|
|
|
# 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_modules), \
|
|
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()
|