zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-17 06:56:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: address review feedback on media metadata extraction and tests

Browse Source
This commit is contained in:
Terry Jia
2026-08-14 21:24:41 -04:00
parent 8969c5cd19
commit c59e25288f
4 changed files with 116 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
app/assets/services/video_metadata.py
View File

@@ -54,8 +54,6 @@ def extract_video_metadata(
elif container.duration is not None:
duration = float(container.duration / av.time_base)
frame_count = stream.frames or None
if frame_count is None and duration is not None and fps is not None:
frame_count = round(duration * fps)
width = stream.codec_context.width
height = stream.codec_context.height

6
comfy_execution/media_enrichment.py
View File

@@ -12,6 +12,8 @@ import logging
import mimetypes
import os
import folder_paths
def resolve_output_entry_path(entry) -> str | None:
"""Resolve a file-type output entry to an absolute path inside its base dir.
@@ -21,12 +23,10 @@ def resolve_output_entry_path(entry) -> str | None:
folder-type base directory (symlinks are resolved before the containment
check), and paths with no file on disk.
"""
import folder_paths
if not isinstance(entry, dict) or "filename" not in entry or "type" not in entry:
return None
filename = entry["filename"]
subfolder = entry.get("subfolder") or ""
subfolder = entry.get("subfolder", "")
if not isinstance(filename, str) or not isinstance(subfolder, str) or not isinstance(entry["type"], str):
return None
base = folder_paths.get_directory_by_type(entry["type"])

66
tests-unit/assets_test/services/test_video_metadata.py
View File

@@ -9,14 +9,18 @@ import pytest
from app.assets.services.media_metadata import extract_media_metadata
from app.assets.services.video_metadata import extract_video_metadata
av = pytest.importorskip("av")
def _make_mp4(
path: Path, width: int = 64, height: int = 48, frames: int = 12, fps: int = 8
def _make_video(
path: Path,
width: int = 64,
height: int = 48,
frames: int = 12,
fps: int = 8,
codec: str = "libx264",
) -> Path:
av = pytest.importorskip("av")
with av.open(str(path), "w") as container:
stream = container.add_stream("libx264", rate=fps)
stream = container.add_stream(codec, rate=fps)
stream.width = width
stream.height = height
stream.pix_fmt = "yuv420p"
@@ -31,6 +35,12 @@ def _make_mp4(
return path
def _make_mp4(
path: Path, width: int = 64, height: int = 48, frames: int = 12, fps: int = 8
) -> Path:
return _make_video(path, width=width, height=height, frames=frames, fps=fps)
class TestExtractVideoMetadata:
def test_extracts_stream_metadata(self, tmp_path: Path):
f = _make_mp4(tmp_path / "clip.mp4", width=64, height=48, frames=12, fps=8)
@@ -79,6 +89,15 @@ class TestExtractVideoMetadata:
assert result is None
def test_webm_duration_from_container_without_frame_count(self, tmp_path: Path):
f = _make_video(tmp_path / "clip.webm", frames=12, fps=8, codec="libvpx-vp9")
result = extract_video_metadata(str(f), mime_type="video/webm")
assert result is not None
assert result["duration"] == pytest.approx(12 / 8, abs=0.2)
assert "frame_count" not in result
class TestExtractMediaMetadata:
def test_dispatches_video_mime_to_video_extractor(self, tmp_path: Path):
@@ -90,7 +109,7 @@ class TestExtractMediaMetadata:
assert result["kind"] == "video"
def test_dispatches_image_mime_to_image_extractor(self, tmp_path: Path):
from PIL import Image
Image = pytest.importorskip("PIL.Image")
f = tmp_path / "img.png"
Image.new("RGB", (32, 16)).save(f, format="PNG")
@@ -136,3 +155,38 @@ class TestIngestStoresVideoMetadata:
assert meta["height"] == 48
assert meta["frame_count"] == 12
assert meta["fps"] == pytest.approx(8.0)
def test_register_existing_asset_backfills_metadata_from_sibling(
self, mock_create_session, temp_dir: Path, session
):
from app.assets.database.models import AssetReference
from app.assets.services.ingest import (
_ingest_file_from_path,
_register_existing_asset,
)
f = _make_mp4(temp_dir / "clip.mp4", width=64, height=48)
_ingest_file_from_path(
abs_path=str(f),
asset_hash="blake3:videosibling",
size_bytes=f.stat().st_size,
mtime_ns=1234567890000000000,
mime_type="video/mp4",
)
result = _register_existing_asset(
asset_hash="blake3:videosibling",
name="copy.mp4",
mime_type="video/mp4",
)
assert result.created is True
session.expire_all()
ref = session.query(AssetReference).filter_by(name="copy.mp4").one()
meta = ref.system_metadata or {}
assert meta["kind"] == "video"
assert meta["width"] == 64
assert meta["height"] == 48
assert meta["frame_count"] == 12
assert meta["fps"] == pytest.approx(8.0)
assert "duration" in meta

93
tests-unit/execution_test/test_media_enrichment.py
View File

@@ -1,17 +1,22 @@
"""Tests for enrich_output_with_media_metadata in comfy_execution/media_enrichment.py."""
import mimetypes
import os
import shutil
import tempfile
import unittest
from unittest.mock import MagicMock, patch
from utils.mime_types import init_mime_types
# 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()
# patched-True isfile makes it try to open files that don't exist. Use the
# project initializer, not mimetypes.init(), which would wipe the custom
# registrations other test modules rely on.
init_mime_types()
# 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")
_DEFAULT_BASE = os.path.join(tempfile.gettempdir(), "media-enrichment-test-base")
_VIDEO_META = {
"kind": "video",
@@ -23,25 +28,25 @@ _VIDEO_META = {
}
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)),
),
}
import comfy_execution.media_enrichment as media_enrichment
def _folder_paths_mock(directory=_DEFAULT_BASE):
return MagicMock(get_directory_by_type=MagicMock(return_value=directory))
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)
extractor_module = MagicMock(extract_media_metadata=extract_mock)
# Only os.path.isfile is patched — abspath/join must run natively so the
# folder_paths is bound at media_enrichment module scope, so it is patched
# as an attribute; the extractor is looked up lazily via sys.modules. 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), \
with patch.object(media_enrichment, "folder_paths", _folder_paths_mock(directory)), \
patch.dict("sys.modules", {"app.assets.services.media_metadata": extractor_module}), \
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
return media_enrichment.enrich_output_with_media_metadata(output_ui), extract_mock
class TestEnrichOutputWithMediaMetadata(unittest.TestCase):
@@ -117,15 +122,11 @@ class TestEnrichOutputWithMediaMetadata(unittest.TestCase):
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), \
# A None sys.modules entry makes the lazy import raise ImportError.
with patch.object(media_enrichment, "folder_paths", _folder_paths_mock()), \
patch.dict("sys.modules", {"app.assets.services.media_metadata": None}), \
patch("os.path.isfile", return_value=True):
import comfy_execution.media_enrichment as mod
result = mod.enrich_output_with_media_metadata(output)
result = media_enrichment.enrich_output_with_media_metadata(output)
self.assertNotIn("metadata", result["images"][0])
def test_extractor_import_failure_beyond_importerror_degrades(self):
@@ -134,16 +135,30 @@ class TestEnrichOutputWithMediaMetadata(unittest.TestCase):
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), \
with patch.object(media_enrichment, "folder_paths", _folder_paths_mock()), \
patch.dict("sys.modules", {"app.assets.services.media_metadata": ExplodingModule()}), \
patch("os.path.isfile", return_value=True):
import comfy_execution.media_enrichment as mod
result = mod.enrich_output_with_media_metadata(output)
result = media_enrichment.enrich_output_with_media_metadata(output)
self.assertNotIn("metadata", result["images"][0])
def test_missing_subfolder_key_defaults_to_base_dir(self):
output = {"images": [{"filename": "clip.mp4", "type": "output"}]}
result, _ = _call(output)
self.assertIn("metadata", result["images"][0])
def test_falsy_non_string_subfolder_skipped(self):
output = {
"images": [
{"filename": "a.mp4", "subfolder": None, "type": "output"},
{"filename": "b.mp4", "subfolder": False, "type": "output"},
{"filename": "c.mp4", "subfolder": 0, "type": "output"},
]
}
result, extract_mock = _call(output)
for entry in result["images"]:
self.assertNotIn("metadata", entry)
extract_mock.assert_not_called()
def test_non_string_fields_skipped(self):
output = {
"images": [
@@ -158,9 +173,6 @@ class TestEnrichOutputWithMediaMetadata(unittest.TestCase):
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")
@@ -168,18 +180,19 @@ class TestEnrichOutputWithMediaMetadata(unittest.TestCase):
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"))
try:
os.symlink(secret, os.path.join(base, "clip.mp4"))
except OSError as e:
self.skipTest(f"cannot create symlinks on this platform: {e}")
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(
with patch.object(media_enrichment, "folder_paths", _folder_paths_mock(base)):
escaped = media_enrichment.resolve_output_entry_path(
{"filename": "clip.mp4", "subfolder": "", "type": "output"})
contained = mod.resolve_output_entry_path(
contained = media_enrichment.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))