Blame for tests/test_h3_guide.py

import os
import sys
import unittest
from unittest.mock import MagicMock

# Add project root to sys.path
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

import asyncio

if 'torch' not in sys.modules:
    sys.modules['torch'] = MagicMock()
if 'numpy' not in sys.modules:
    sys.modules['numpy'] = MagicMock()
if 'gradio' not in sys.modules:
    sys.modules['gradio'] = MagicMock()
if 'spaces' not in sys.modules:
    mock_spaces = MagicMock()
    mock_spaces.GPU = lambda fn: fn
    sys.modules['spaces'] = mock_spaces
if 'huggingface_hub' not in sys.modules:
    mock_hf = MagicMock()
    sys.modules['huggingface_hub'] = mock_hf
    sys.modules['huggingface_hub.constants'] = MagicMock()
if 'PIL' not in sys.modules:
    mock_pil = MagicMock()
    class MockImage:
        pass
    mock_pil.Image = MagicMock()
    mock_pil.Image.Image = MockImage
    sys.modules['PIL'] = mock_pil
    sys.modules['PIL.Image'] = mock_pil.Image

# Mock ComfyUI environment if not installed locally
if 'execution' not in sys.modules:
    sys.modules['execution'] = MagicMock()
if 'server' not in sys.modules:
    sys.modules['server'] = MagicMock()
if 'nodes' not in sys.modules:
    mock_nodes = MagicMock()
    async def async_noop():
        pass
    mock_nodes.init_extra_nodes = async_noop
    dummy_cls = MagicMock()
    mock_nodes.NODE_CLASS_MAPPINGS = {
        'CLIPTextEncode': dummy_cls,
        'CLIPTextEncodeSDXL': dummy_cls,
        'LoraLoader': dummy_cls,
        'CLIPSetLastLayer': dummy_cls,
        'CheckpointLoaderSimple': dummy_cls,
    }
    mock_nodes.CheckpointLoaderSimple = dummy_cls
    mock_nodes.LoraLoader = dummy_cls
    sys.modules['nodes'] = mock_nodes

# Helper to create dummy ComfyUI node class
def make_mock_node(name, required_inputs=None, optional_inputs=None):
    req = required_inputs or {}
    opt = optional_inputs or {}
    class MockNode:
        NODE_NAME = name
        @classmethod
        def INPUT_TYPES(cls):
            return {"required": req, "optional": opt}
    return MockNode

from comfy_integration import nodes as ci_nodes

# Register mock nodes for testing
TEST_NODE_MAPPINGS = {
    "UNETLoader": make_mock_node("UNETLoader", {"unet_name": (["model.safetensors"],), "weight_dtype": (["default"],)}),
    "CLIPLoader": make_mock_node("CLIPLoader", {"clip_name": (["clip.safetensors"],), "type": (["minimax"],)}),
    "VAELoader": make_mock_node("VAELoader", {"vae_name": (["vae.safetensors"],)}),
    "MiniMaxH3ReferenceToVideo": make_mock_node("MiniMaxH3ReferenceToVideo", {
        "clip": ("CLIP",), "vae": ("VAE",), "audio_vae": ("VAE",),
        "prompt": ("STRING", {"default": ""}), "width": ("INT", {"default": 960}),
        "height": ("INT", {"default": 544}), "length": ("INT", {"default": 73})
    }),
    "BasicGuider": make_mock_node("BasicGuider", {"model": ("MODEL",), "conditioning": ("CONDITIONING",)}),
    "KSamplerSelect": make_mock_node("KSamplerSelect", {"sampler_name": (["res_multistep"], {"default": "res_multistep"})}),
    "BasicScheduler": make_mock_node("BasicScheduler", {"model": ("MODEL",), "scheduler": (["simple"], {"default": "simple"}), "steps": ("INT", {"default": 20})}),
    "RandomNoise": make_mock_node("RandomNoise", {"noise_seed": ("INT", {"default": 0})}),
    "SamplerCustomAdvanced": make_mock_node("SamplerCustomAdvanced", {
        "noise": ("NOISE",), "guider": ("GUIDER",), "sampler": ("SAMPLER",),
        "sigmas": ("SIGMAS",), "latent_image": ("LATENT",)
    }),
    "VAEDecode": make_mock_node("VAEDecode", {"samples": ("LATENT",), "vae": ("VAE",)}),
    "VAEDecodeAudio": make_mock_node("VAEDecodeAudio", {"samples": ("LATENT",), "vae": ("VAE",)}),
    "CreateVideo": make_mock_node("CreateVideo", {"images": ("IMAGE",), "audio": ("AUDIO",), "fps": ("INT", {"default": 24})}),
    "SaveVideo": make_mock_node("SaveVideo", {"video": ("VIDEO",), "filename_prefix": ("STRING", {"default": "video/test"})}),
    "LoadImage": make_mock_node("LoadImage", {"image": ("STRING", {"default": ""})}),
    "LoadVideo": make_mock_node("LoadVideo", {"file": ("STRING", {"default": ""})}),
    "GetVideoComponents": make_mock_node("GetVideoComponents", {"video": ("VIDEO",)}),
    "LoadAudio": make_mock_node("LoadAudio", {"audio": ("STRING", {"default": ""})}),
    "MiniMaxH3AddGuide": make_mock_node("MiniMaxH3AddGuide", {
        "positive": ("CONDITIONING",),
        "latent": ("LATENT",),
        "frame_idx": ("INT", {"default": 0}),
    }, {
        "vae": ("VAE",),
        "audio_vae": ("VAE",),
        "image": ("IMAGE",),
        "audio": ("AUDIO",),
    }),
    "LoraLoader": make_mock_node("LoraLoader", {"model": ("MODEL",), "clip": ("CLIP",), "lora_name": ("STRING",)}),
    "MiniMaxH3ImageToVideo": make_mock_node("MiniMaxH3ImageToVideo", {
        "clip": ("CLIP",), "vae": ("VAE",),
        "prompt": ("STRING", {"default": ""}), "width": ("INT", {"default": 1344}),
        "height": ("INT", {"default": 768}), "length": ("INT", {"default": 73})
    }, {
        "first_frame": ("IMAGE",), "last_frame": ("IMAGE",)
    }),
    "ImageScale": make_mock_node("ImageScale", {
        "image": ("IMAGE",), "width": ("INT", {"default": 512}), "height": ("INT", {"default": 512}),
        "upscale_method": (["nearest-exact"],), "crop": (["center"],)
    }),
}

ci_nodes.NODE_CLASS_MAPPINGS.update(TEST_NODE_MAPPINGS)

from chain_injectors import discover_injectors, get_registered_features
from utils.app_utils import process_h3_guide_inputs
from core.workflow_assembler import WorkflowAssembler


class TestH3GuideIntegration(unittest.TestCase):
    def setUp(self):
        self.recipe_path = os.path.join(PROJECT_ROOT, "ui", "shared", "h3_ref2va_recipe.yaml")

    def test_injector_discovery(self):
        injectors = discover_injectors()
        self.assertIn("dynamic_h3_guide_chains", injectors)
        self.assertTrue(callable(injectors["dynamic_h3_guide_chains"]))

    def test_registered_features(self):
        import yaml
        features = get_registered_features()
        self.assertIn("h3_guide", features)
        self.assertEqual(features["h3_guide"]["chain_type"], "dynamic_h3_guide_chains")

        with open(os.path.join(PROJECT_ROOT, "yaml", "chain_features.yaml"), "r", encoding="utf-8") as f:
            chain_features = yaml.safe_load(f)
        self.assertIn("h3_guide", chain_features)
        schema = chain_features["h3_guide"]["parameters_schema"]
        self.assertIn("frame_idx", schema["properties"])
        self.assertIn("time_seconds", schema["properties"])
        self.assertIn("image", schema["properties"])
        self.assertIn("video", schema["properties"])
        self.assertIn("audio", schema["properties"])

    def test_process_h3_guide_inputs_direct_list(self):
        raw_guides = [
            {"image": "path/to/img1.png", "time_seconds": 1.5},
            {"image": "path/to/img2.png", "frame_idx": 10},
        ]
        ui_values = {"h3_guides": raw_guides}
        processed = process_h3_guide_inputs(ui_values)
        self.assertEqual(len(processed), 2)
        # Should be sorted by frame_idx: 10 comes before round(1.5 * 24) = 36
        self.assertEqual(processed[0]["frame_idx"], 10)
        self.assertEqual(processed[1]["frame_idx"], 36)
        self.assertEqual(processed[0]["image"], "img2.png")
        self.assertEqual(processed[1]["image"], "img1.png")

    def test_process_h3_guide_inputs_ui_prefix(self):
        ui_values = {
            "h3_guide_images_test": ["img_a.png", "img_b.png"],
            "h3_guide_videos_test": [None, None],
            "h3_guide_audios_test": [None, None],
            "h3_guide_times_test": [2.0, 0.5],
            "h3_guide_frames_test": [None, None],
        }
        processed = process_h3_guide_inputs(ui_values, prefix="test")
        self.assertEqual(len(processed), 2)
        # Sorted: 0.5s (12 frames) then 2.0s (48 frames)
        self.assertEqual(processed[0]["frame_idx"], 12)
        self.assertEqual(processed[1]["frame_idx"], 48)
        self.assertEqual(processed[0]["image"], "img_b.png")
        self.assertEqual(processed[1]["image"], "img_a.png")

    def test_workflow_assembler_without_guides(self):
        assembler = WorkflowAssembler(self.recipe_path)
        ui_values = {
            "prompt": "test prompt",
            "width": 960,
            "height": 544,
            "length": 73,
            "h3_guides": []
        }
        workflow = assembler.assemble(ui_values)
        # BasicGuider conditioning should connect directly to minimax_h3:0
        guider_id = assembler.node_map["basic_guider"]
        minimax_id = assembler.node_map["minimax_h3"]
        cond_input = workflow[guider_id]["inputs"]["conditioning"]
        self.assertEqual(cond_input, [minimax_id, 0])

    def test_workflow_assembler_with_single_guide(self):
        assembler = WorkflowAssembler(self.recipe_path)
        ui_values = {
            "prompt": "test prompt",
            "width": 960,
            "height": 544,
            "length": 73,
            "h3_guides": [
                {"image": "test_guide.png", "frame_idx": 24}
            ]
        }
        workflow = assembler.assemble(ui_values)
        guider_id = assembler.node_map["basic_guider"]
        minimax_id = assembler.node_map["minimax_h3"]
        vae_id = assembler.node_map["video_vae_loader"]
        audio_vae_id = assembler.node_map["audio_vae_loader"]

        # Find MiniMaxH3AddGuide node in workflow
        guide_nodes = [
            (nid, n) for nid, n in workflow.items()
            if n.get("class_type") == "MiniMaxH3AddGuide"
        ]
        self.assertEqual(len(guide_nodes), 1)
        guide_id, guide_data = guide_nodes[0]

        # Verify guide inputs
        self.assertEqual(guide_data["inputs"]["frame_idx"], 24)
        self.assertEqual(guide_data["inputs"]["positive"], [minimax_id, 0])
        self.assertEqual(guide_data["inputs"]["latent"], [minimax_id, 1])
        self.assertEqual(guide_data["inputs"]["vae"], [vae_id, 0])
        self.assertEqual(guide_data["inputs"]["audio_vae"], [audio_vae_id, 0])

        # Verify basic_guider points to the guide output
        self.assertEqual(workflow[guider_id]["inputs"]["conditioning"], [guide_id, 0])

    def test_workflow_assembler_with_chained_guides(self):
        assembler = WorkflowAssembler(self.recipe_path)
        ui_values = {
            "prompt": "test prompt",
            "width": 960,
            "height": 544,
            "length": 73,
            "h3_guides": [
                {"image": "guide1.png", "frame_idx": 10},
                {"image": "guide2.png", "frame_idx": 36},
                {"image": "guide3.png", "frame_idx": 72}
            ]
        }
        workflow = assembler.assemble(ui_values)
        guider_id = assembler.node_map["basic_guider"]
        minimax_id = assembler.node_map["minimax_h3"]

        guide_nodes = [
            (nid, n) for nid, n in workflow.items()
            if n.get("class_type") == "MiniMaxH3AddGuide"
        ]
        self.assertEqual(len(guide_nodes), 3)

        # Sort by frame_idx
        guide_nodes_sorted = sorted(guide_nodes, key=lambda x: x[1]["inputs"]["frame_idx"])
        g1_id, g1 = guide_nodes_sorted[0]
        g2_id, g2 = guide_nodes_sorted[1]
        g3_id, g3 = guide_nodes_sorted[2]

        self.assertEqual(g1["inputs"]["positive"], [minimax_id, 0])
        self.assertEqual(g2["inputs"]["positive"], [g1_id, 0])
        self.assertEqual(g3["inputs"]["positive"], [g2_id, 0])

        # BasicGuider connects to last guide in chain
        self.assertEqual(workflow[guider_id]["inputs"]["conditioning"], [g3_id, 0])

    def test_workflow_assembler_with_video_and_audio_guides(self):
        assembler = WorkflowAssembler(self.recipe_path)
        ui_values = {
            "prompt": "test prompt",
            "width": 960,
            "height": 544,
            "length": 73,
            "h3_guides": [
                {"video": "clip.mp4", "frame_idx": 24},
                {"audio": "audio.mp3", "frame_idx": 48}
            ]
        }
        workflow = assembler.assemble(ui_values)

        guide_nodes = [
            (nid, n) for nid, n in workflow.items()
            if n.get("class_type") == "MiniMaxH3AddGuide"
        ]
        self.assertEqual(len(guide_nodes), 2)

        # Check LoadVideo + GetVideoComponents created
        load_vids = [n for n in workflow.values() if n.get("class_type") == "LoadVideo"]
        get_comps = [n for n in workflow.values() if n.get("class_type") == "GetVideoComponents"]
        load_auds = [n for n in workflow.values() if n.get("class_type") == "LoadAudio"]

        self.assertEqual(len(load_vids), 1)
        self.assertEqual(len(get_comps), 1)
        self.assertEqual(len(load_auds), 1)

    def test_h3_guide_ui_configuration(self):
        from ui.shared.ui_components import create_h3_guide_ui
        import gradio as gr

        gr.Slider.reset_mock()
        gr.Number.reset_mock()

        comps = create_h3_guide_ui("test_ui")
        times = comps.get("h3_guide_times_test_ui", [])
        frames = comps.get("h3_guide_frames_test_ui", [])

        # Step must be 0.5 for all slider calls
        self.assertGreater(len(times), 0)
        self.assertGreater(len(gr.Slider.call_args_list), 0)
        for call in gr.Slider.call_args_list:
            self.assertEqual(call.kwargs.get("step"), 0.5)

        # No Frame Index components displayed
        self.assertEqual(len(frames), 0)
        for call in gr.Number.call_args_list:
            label = call.kwargs.get("label", "")
            self.assertNotIn("Frame Index", label)
            self.assertNotIn("24 fps", label.lower())

    def test_h3_fl2va_recipe_guide_integration(self):
        fl2va_recipe_path = os.path.join(PROJECT_ROOT, "ui", "shared", "h3_fl2va_recipe.yaml")
        assembler = WorkflowAssembler(fl2va_recipe_path)
        ui_values = {
            "prompt": "sunset test",
            "width": 1344,
            "height": 768,
            "length": 73,
            "h3_guides": [
                {"image": "intermediate_keyframe.png", "frame_idx": 36}
            ]
        }
        workflow = assembler.assemble(ui_values)
        guider_id = assembler.node_map["basic_guider"]
        minimax_id = assembler.node_map["minimax_h3"]

        guide_nodes = [
            (nid, n) for nid, n in workflow.items()
            if n.get("class_type") == "MiniMaxH3AddGuide"
        ]
        self.assertEqual(len(guide_nodes), 1)
        guide_id, guide_data = guide_nodes[0]
        self.assertEqual(guide_data["inputs"]["frame_idx"], 36)
        self.assertEqual(guide_data["inputs"]["positive"], [minimax_id, 0])
        self.assertEqual(workflow[guider_id]["inputs"]["conditioning"], [guide_id, 0])

    def test_h3_fl2va_ui_integration(self):
        from ui.shared import h3_fl2va_ui

        comps = h3_fl2va_ui.create_ui()
        self.assertIn("h3_guide_times_h3_fl2va", comps)
        self.assertIn("h3_guide_images_h3_fl2va", comps)
        self.assertIn("h3_guide_videos_h3_fl2va", comps)
        self.assertIn("h3_guide_audios_h3_fl2va", comps)

    def test_unified_h3_ui_integration(self):
        from ui.h3_ui import create_ui
        comps = create_ui()
        self.assertIn("task", comps)
        self.assertIn("first_frame", comps)
        self.assertIn("last_frame", comps)
        self.assertIn("h3_guide_times_h3", comps)
        self.assertIn("h3_guide_images_h3", comps)
        self.assertIn("h3_guide_videos_h3", comps)
        self.assertIn("h3_guide_audios_h3", comps)


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