221 lines
9.5 KiB
Python
221 lines
9.5 KiB
Python
"""Tests for enrich_output_with_media_metadata in comfy_execution/media_enrichment.py."""
|
|
import mimetypes
|
|
import os
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# Initialize the mimetypes registry before any test patches os.path.isfile —
|
|
# its lazy init consults os.path.isfile to pick candidate files, and a
|
|
# patched-True isfile makes it try to open files that don't exist.
|
|
mimetypes.init()
|
|
|
|
# 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(), "media-enrichment-test-base")
|
|
|
|
_VIDEO_META = {
|
|
"kind": "video",
|
|
"width": 1280,
|
|
"height": 720,
|
|
"duration": 4.0,
|
|
"fps": 24.0,
|
|
"frame_count": 96,
|
|
}
|
|
|
|
|
|
def _mocked_modules(*, extract=None, directory=_DEFAULT_BASE):
|
|
return {
|
|
"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=directory)),
|
|
"app.assets.services.media_metadata": MagicMock(
|
|
extract_media_metadata=extract or MagicMock(return_value=dict(_VIDEO_META)),
|
|
),
|
|
}
|
|
|
|
|
|
def _call(output_ui, *, extract=None, file_exists=True, directory=_DEFAULT_BASE):
|
|
extract_mock = extract or MagicMock(return_value=dict(_VIDEO_META))
|
|
mocked = _mocked_modules(extract=extract_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 comfy_execution.media_enrichment as mod
|
|
return mod.enrich_output_with_media_metadata(output_ui), extract_mock
|
|
|
|
|
|
class TestEnrichOutputWithMediaMetadata(unittest.TestCase):
|
|
|
|
def test_attaches_metadata_object(self):
|
|
output = {"images": [{"filename": "clip.mp4", "subfolder": "", "type": "output"}]}
|
|
result, _ = _call(output)
|
|
self.assertEqual(result["images"][0]["metadata"], _VIDEO_META)
|
|
|
|
def test_passes_guessed_mime_type(self):
|
|
output = {"images": [{"filename": "clip.mp4", "subfolder": "", "type": "output"}]}
|
|
_, extract_mock = _call(output)
|
|
_, kwargs = extract_mock.call_args
|
|
self.assertEqual(kwargs["mime_type"], "video/mp4")
|
|
|
|
def test_original_entry_not_mutated(self):
|
|
orig = {"filename": "clip.mp4", "subfolder": "", "type": "output"}
|
|
_call({"images": [orig]})
|
|
self.assertNotIn("metadata", orig)
|
|
|
|
def test_non_media_extractor_none_leaves_entry_unchanged(self):
|
|
extract = MagicMock(return_value=None)
|
|
output = {"latent": [{"filename": "a.latent", "subfolder": "", "type": "output"}]}
|
|
result, _ = _call(output, extract=extract)
|
|
self.assertNotIn("metadata", result["latent"][0])
|
|
|
|
def test_existing_metadata_key_untouched(self):
|
|
entry = {"filename": "clip.mp4", "subfolder": "", "type": "output", "metadata": {"kind": "other"}}
|
|
result, extract_mock = _call({"images": [entry]})
|
|
self.assertEqual(result["images"][0]["metadata"], {"kind": "other"})
|
|
extract_mock.assert_not_called()
|
|
|
|
def test_non_list_value_passed_through(self):
|
|
result, _ = _call({"text": "hello"})
|
|
self.assertEqual(result["text"], "hello")
|
|
|
|
def test_none_entry_in_list_unchanged(self):
|
|
output = {"images": [None, {"filename": "a.mp4", "subfolder": "", "type": "output"}]}
|
|
result, _ = _call(output)
|
|
self.assertIsNone(result["images"][0])
|
|
self.assertIn("metadata", result["images"][1])
|
|
|
|
def test_entry_without_filename_unchanged(self):
|
|
output = {"latent": [{"subfolder": "", "type": "output"}]}
|
|
result, extract_mock = _call(output)
|
|
self.assertNotIn("metadata", result["latent"][0])
|
|
extract_mock.assert_not_called()
|
|
|
|
def test_file_not_on_disk_unchanged(self):
|
|
output = {"images": [{"filename": "missing.mp4", "subfolder": "", "type": "output"}]}
|
|
result, extract_mock = _call(output, file_exists=False)
|
|
self.assertNotIn("metadata", result["images"][0])
|
|
extract_mock.assert_not_called()
|
|
|
|
def test_unknown_type_directory_unchanged(self):
|
|
output = {"images": [{"filename": "a.mp4", "subfolder": "", "type": "unknown"}]}
|
|
result, extract_mock = _call(output, directory=None)
|
|
self.assertNotIn("metadata", result["images"][0])
|
|
extract_mock.assert_not_called()
|
|
|
|
def test_path_traversal_subfolder_skipped(self):
|
|
output = {"images": [{"filename": "passwd", "subfolder": "../../etc", "type": "output"}]}
|
|
result, extract_mock = _call(output)
|
|
self.assertNotIn("metadata", result["images"][0])
|
|
extract_mock.assert_not_called()
|
|
|
|
def test_absolute_filename_skipped(self):
|
|
absolute_filename = os.path.abspath(os.sep + "etc" + os.sep + "passwd")
|
|
output = {"images": [{"filename": absolute_filename, "subfolder": "", "type": "output"}]}
|
|
result, extract_mock = _call(output)
|
|
self.assertNotIn("metadata", result["images"][0])
|
|
extract_mock.assert_not_called()
|
|
|
|
def test_extractor_unavailable_returns_unchanged(self):
|
|
output = {"images": [{"filename": "a.mp4", "subfolder": "", "type": "output"}]}
|
|
mocked = {
|
|
"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)),
|
|
# A None sys.modules entry makes the lazy import raise ImportError.
|
|
"app.assets.services.media_metadata": None,
|
|
}
|
|
with patch.dict("sys.modules", mocked), \
|
|
patch("os.path.isfile", return_value=True):
|
|
import comfy_execution.media_enrichment as mod
|
|
result = mod.enrich_output_with_media_metadata(output)
|
|
self.assertNotIn("metadata", result["images"][0])
|
|
|
|
def test_extractor_import_failure_beyond_importerror_degrades(self):
|
|
class ExplodingModule:
|
|
def __getattr__(self, name):
|
|
raise RuntimeError("dependency init failed")
|
|
|
|
output = {"images": [{"filename": "a.mp4", "subfolder": "", "type": "output"}]}
|
|
mocked = {
|
|
"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)),
|
|
"app.assets.services.media_metadata": ExplodingModule(),
|
|
}
|
|
with patch.dict("sys.modules", mocked), \
|
|
patch("os.path.isfile", return_value=True):
|
|
import comfy_execution.media_enrichment as mod
|
|
result = mod.enrich_output_with_media_metadata(output)
|
|
self.assertNotIn("metadata", result["images"][0])
|
|
|
|
def test_non_string_fields_skipped(self):
|
|
output = {
|
|
"images": [
|
|
{"filename": 42, "subfolder": "", "type": "output"},
|
|
{"filename": "a.mp4", "subfolder": ["nested"], "type": "output"},
|
|
{"filename": "b.mp4", "subfolder": "", "type": 7},
|
|
]
|
|
}
|
|
result, extract_mock = _call(output)
|
|
for entry in result["images"]:
|
|
self.assertNotIn("metadata", entry)
|
|
extract_mock.assert_not_called()
|
|
|
|
def test_symlink_escape_rejected_real_fs(self):
|
|
import shutil
|
|
import tempfile
|
|
|
|
root = tempfile.mkdtemp(prefix="media-enrichment-symlink-")
|
|
try:
|
|
base = os.path.join(root, "output")
|
|
os.makedirs(base)
|
|
secret = os.path.join(root, "secret.txt")
|
|
with open(secret, "w") as f:
|
|
f.write("outside")
|
|
os.symlink(secret, os.path.join(base, "clip.mp4"))
|
|
inside = os.path.join(base, "real.mp4")
|
|
with open(inside, "w") as f:
|
|
f.write("inside")
|
|
|
|
mocked = {"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=base))}
|
|
# No isfile patching — this test runs against the real filesystem.
|
|
with patch.dict("sys.modules", mocked):
|
|
import comfy_execution.media_enrichment as mod
|
|
escaped = mod.resolve_output_entry_path(
|
|
{"filename": "clip.mp4", "subfolder": "", "type": "output"})
|
|
contained = mod.resolve_output_entry_path(
|
|
{"filename": "real.mp4", "subfolder": "", "type": "output"})
|
|
self.assertIsNone(escaped, "symlink pointing outside the base must be rejected")
|
|
self.assertEqual(contained, os.path.realpath(inside))
|
|
finally:
|
|
shutil.rmtree(root, ignore_errors=True)
|
|
|
|
def test_extractor_error_does_not_block_sibling_entries(self):
|
|
call_count = [0]
|
|
|
|
def extract_side_effect(abs_path, mime_type=None):
|
|
call_count[0] += 1
|
|
if call_count[0] == 1:
|
|
raise RuntimeError("boom")
|
|
return dict(_VIDEO_META)
|
|
|
|
extract = MagicMock(side_effect=extract_side_effect)
|
|
output = {
|
|
"images": [
|
|
{"filename": "bad.mp4", "subfolder": "", "type": "output"},
|
|
{"filename": "good.mp4", "subfolder": "", "type": "output"},
|
|
]
|
|
}
|
|
result, _ = _call(output, extract=extract)
|
|
self.assertNotIn("metadata", result["images"][0])
|
|
self.assertEqual(result["images"][1]["metadata"], _VIDEO_META)
|
|
|
|
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("metadata", result["images"][0])
|
|
self.assertIn("metadata", result["videos"][0])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|