File size: 8,219 Bytes
04eb75a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | """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()
|