ComfyAssets-ComfyUI_PromptManager/tests/test_utils_video_metadata.py

"""Tests for utils/video_metadata.py: prompt metadata from video containers.

ComfyUI's video savers store ``{"prompt": ..., "workflow": ...}`` as JSON in
the container's ``comment`` tag. ffprobe reads it; when ffprobe is missing
videos are skipped without noise.
"""

import json
import os
import shutil
import subprocess
import sys
import tempfile
import unittest
from unittest import mock

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils import video_metadata  # noqa: E402
from utils.video_metadata import (  # noqa: E402
    VIDEO_EXTENSIONS,
    is_video_path,
    metadata_from_tags,
    read_video_metadata,
)

HAVE_FFMPEG = bool(shutil.which("ffmpeg") and shutil.which("ffprobe"))
GRAPH = json.dumps({"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}})


def write_mp4(path, comment):
    """Write a tiny black video carrying ``comment`` in its container tags."""
    subprocess.run(
        [
            "ffmpeg",
            "-v",
            "quiet",
            "-f",
            "lavfi",
            "-i",
            "color=c=black:s=16x16:d=0.2",
            "-metadata",
            f"comment={comment}",
            "-y",
            str(path),
        ],
        check=True,
        timeout=60,
    )


class TestMetadataFromTags(unittest.TestCase):

    def test_prompt_string_and_workflow_dict_become_text_chunks(self):
        tags = {"comment": json.dumps({"prompt": GRAPH, "workflow": {"nodes": []}})}
        meta = metadata_from_tags(tags)
        self.assertEqual(meta["prompt"], GRAPH)
        self.assertEqual(json.loads(meta["workflow"]), {"nodes": []})

    def test_missing_or_non_json_comment_is_empty(self):
        self.assertEqual(metadata_from_tags({}), {})
        self.assertEqual(metadata_from_tags({"comment": "just a note"}), {})
        self.assertEqual(metadata_from_tags({"comment": "[1, 2]"}), {})
        self.assertEqual(metadata_from_tags({"comment": json.dumps({"other": 1})}), {})
        self.assertEqual(metadata_from_tags({"comment": 7}), {})

    def test_description_tag_is_a_fallback(self):
        tags = {"description": json.dumps({"prompt": GRAPH})}
        self.assertEqual(metadata_from_tags(tags), {"prompt": GRAPH})


class TestIsVideoPath(unittest.TestCase):

    def test_extensions(self):
        for ext in VIDEO_EXTENSIONS:
            self.assertTrue(is_video_path(f"clip{ext.upper()}"), ext)
        self.assertFalse(is_video_path("image.png"))
        self.assertFalse(is_video_path("noext"))


class TestReadVideoMetadataWithoutFfprobe(unittest.TestCase):

    def setUp(self):
        video_metadata.ffprobe_path.cache_clear()
        self.addCleanup(video_metadata.ffprobe_path.cache_clear)

    def test_missing_ffprobe_returns_empty_without_spawning(self):
        with mock.patch.object(video_metadata.shutil, "which", return_value=None):
            with mock.patch.object(video_metadata.subprocess, "run") as run:
                self.assertEqual(read_video_metadata("clip.mp4"), {})
        run.assert_not_called()

    def test_missing_ffprobe_is_mentioned_once(self):
        video_metadata.reset_missing_notice()
        with mock.patch.object(video_metadata.shutil, "which", return_value=None):
            with self.assertLogs("prompt_manager.video", level="INFO") as cm:
                read_video_metadata("a.mp4")
                read_video_metadata("b.mp4")
        self.assertEqual(len(cm.output), 1)
        self.assertIn("ffprobe", cm.output[0])

    def _run_returning(self, **attrs):
        return mock.patch.object(
            video_metadata.subprocess,
            "run",
            return_value=mock.Mock(**{"returncode": 0, "stdout": "{}", **attrs}),
        )

    def test_ffprobe_failure_timeout_and_garbage_are_empty(self):
        with mock.patch.object(
            video_metadata.shutil, "which", return_value="/bin/ffprobe"
        ):
            with self._run_returning(returncode=1):
                self.assertEqual(read_video_metadata("clip.mp4"), {})
            with self._run_returning(stdout="not json"):
                self.assertEqual(read_video_metadata("clip.mp4"), {})
            with mock.patch.object(
                video_metadata.subprocess,
                "run",
                side_effect=subprocess.TimeoutExpired("ffprobe", 1),
            ):
                self.assertEqual(read_video_metadata("clip.mp4"), {})
            with mock.patch.object(
                video_metadata.subprocess, "run", side_effect=OSError
            ):
                self.assertEqual(read_video_metadata("clip.mp4"), {})

    def test_ffprobe_output_is_parsed(self):
        payload = {"format": {"tags": {"comment": json.dumps({"prompt": GRAPH})}}}
        with mock.patch.object(
            video_metadata.shutil, "which", return_value="/bin/ffprobe"
        ):
            with self._run_returning(stdout=json.dumps(payload)) as run:
                self.assertEqual(read_video_metadata("clip.mp4"), {"prompt": GRAPH})
        cmd = run.call_args.args[0]
        self.assertEqual(cmd[0], "/bin/ffprobe")
        self.assertEqual(cmd[-1], "clip.mp4")
        self.assertIn("format_tags", cmd)


@unittest.skipUnless(HAVE_FFMPEG, "ffmpeg/ffprobe not installed")
class TestReadVideoMetadataWithRealFfprobe(unittest.TestCase):

    def setUp(self):
        video_metadata.ffprobe_path.cache_clear()
        self.tmpdir = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, self.tmpdir, True)

    def test_round_trip_through_a_real_mp4(self):
        path = os.path.join(self.tmpdir, "clip.mp4")
        write_mp4(path, json.dumps({"prompt": GRAPH, "workflow": {"v": 1}}))

        meta = read_video_metadata(path)

        self.assertEqual(meta["prompt"], GRAPH)
        self.assertEqual(json.loads(meta["workflow"]), {"v": 1})

    def test_video_without_a_comment_is_empty(self):
        path = os.path.join(self.tmpdir, "plain.mp4")
        write_mp4(path, "")
        self.assertEqual(read_video_metadata(path), {})


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