From 43ca97accbfbb87dcb363e7c68a5f946c9802d94 Mon Sep 17 00:00:00 2001 From: Meisoftcoltd <38667978+Meisoftcoltd@users.noreply.github.com> Date: Wed, 24 Jun 2026 23:03:31 +0000 Subject: [PATCH] feat: add DynamicSceneDirector to loop.py and cleanup tests - Injected DynamicSceneDirector node logic at the end of loop.py - Registered node name in __init__.py - Removed tests directory and test files to clean up the repository for production --- __init__.py | 1 + loop.py | 92 ++++++++++++++++++ run_tests.sh | 19 ---- test_loop_syntax.py | 24 ----- tests/sandbox/input/valid.csv | 2 - tests/sandbox/secret.txt | 2 - tests/test_install_fix.py | 34 ------- tests/test_loop_calculator.py | 170 ---------------------------------- tests/test_security_fix.py | 26 ------ tests/test_ssrf_fix.py | 46 --------- tests/test_switch_node.py | 44 --------- tests/test_tools.py | 50 ---------- tests/test_video_audio.py | 162 -------------------------------- tests/test_vram_node.py | 55 ----------- tests/verify_fix.py | 118 ----------------------- 15 files changed, 93 insertions(+), 752 deletions(-) delete mode 100755 run_tests.sh delete mode 100644 test_loop_syntax.py delete mode 100644 tests/sandbox/input/valid.csv delete mode 100644 tests/sandbox/secret.txt delete mode 100644 tests/test_install_fix.py delete mode 100644 tests/test_loop_calculator.py delete mode 100644 tests/test_security_fix.py delete mode 100644 tests/test_ssrf_fix.py delete mode 100644 tests/test_switch_node.py delete mode 100644 tests/test_tools.py delete mode 100644 tests/test_video_audio.py delete mode 100644 tests/test_vram_node.py delete mode 100644 tests/verify_fix.py diff --git a/__init__.py b/__init__.py index 70013f5..df187da 100644 --- a/__init__.py +++ b/__init__.py @@ -33,6 +33,7 @@ NODE_DISPLAY_NAME_MAPPINGS.update({ "SequentialLoopStart": "🏁 Loop Start (Index)", "SequentialLoopTrigger": "🚀 Loop Trigger (Auto-Queue)", "SequentialAudioBatchLoader": "🎵 Sequential Audio Batch Loader", + "DynamicSceneDirector": "🎬 Dynamic Scene Director", # Tools "AutoFPSLimiter": "⏱️ Auto FPS Limiter", "VAESafeFramePadder": "🛡️ VAE Safe Frame Padder", diff --git a/loop.py b/loop.py index 1f648cb..a8d6b42 100644 --- a/loop.py +++ b/loop.py @@ -567,3 +567,95 @@ class SequentialAudioBatchLoader: _log(f" -> ✅ Lote configurado: step_by_chunk={loop.global_step_by_chunk}, acumulado={loop.global_accumulated_frames}/{loop.global_source_frame_count}, is_final={loop.global_is_final_chunk}, has_more_batches={loop.global_has_more_batches}") return (audio_dict, current_file_name, safe_index, total_files, "\n".join(log_output)) + +import folder_paths +import math + +@register_node +class DynamicSceneDirector: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "agent_json": ("STRING", {"multiline": True, "default": '{"scenes": [{"scene_id": 1, "duration_seconds": 5.0, "flux_prompt": "A cinematic wide shot of a cyberpunk city...", "wan_prompt": "Camera pans slowly to the right, neon lights flickering."}]}'}), + "audio_filename": ("STRING", {"default": "audio_track"}), + "current_loop_index": ("INT", {"default": 0, "min": 0, "max": 10000}), + "fps": ("INT", {"default": 24, "min": 8, "max": 60}), + } + } + + RETURN_TYPES = ("BOOLEAN", "STRING", "STRING", "INT", "STRING") + RETURN_NAMES = ("is_flux_phase", "flux_prompt", "wan_prompt", "chunk_frames", "image_load_path") + FUNCTION = "direct_scene" + CATEGORY = "🔁 Sequential Batcher/Director" + + def direct_scene(self, agent_json, audio_filename, current_loop_index, fps): + log_output = [] + def _log(msg): + print(msg); log_output.append(str(msg)) + + _log(f"\n{'='*50}") + _log(f"🎬 [Director] Evaluando Ciclo: {current_loop_index}") + + # 1. Parsear el JSON del Agente + try: + data = json.loads(agent_json) + scenes = data.get("scenes", []) + except Exception as e: + raise ValueError(f"❌ Error en el JSON del Agente: {e}") + + total_scenes = len(scenes) + if total_scenes == 0: + raise ValueError("❌ El JSON no contiene escenas válidas.") + + # 2. Crear el directorio blindado para este proyecto + base_output = folder_paths.get_output_directory() + safe_audio_name = "".join(c for c in audio_filename if c.isalnum() or c in " _-").strip() + if not safe_audio_name: + safe_audio_name = "proyecto_generico" + + keyframes_dir = os.path.join(base_output, f"{safe_audio_name}_Keyframes") + os.makedirs(keyframes_dir, exist_ok=True) + + # 3. LÓGICA DE FASES (FLUX vs WANVIDEO) + if current_loop_index < total_scenes: + # FASE 1: PRE-PRODUCCIÓN (FLUX) + scene = scenes[current_loop_index] + is_flux_phase = True + flux_prompt = scene.get("flux_prompt", "") + wan_prompt = "" + chunk_frames = 0 + + # Nombre de guardado para la imagen de FLUX + image_path = os.path.join(keyframes_dir, f"scene_{scene.get('scene_id', current_loop_index)}.png") + + _log(f" -> 📸 FASE 1 (Pre-Producción): Generando Keyframe para Escena {scene.get('scene_id', current_loop_index)}") + _log(f" -> 💾 Destino del PNG: {image_path}") + + else: + # FASE 2: RENDERIZADO (WANVIDEO) + is_flux_phase = False + video_loop_index = current_loop_index - total_scenes + + # Límite de seguridad + if video_loop_index >= total_scenes: + video_loop_index = total_scenes - 1 + + scene = scenes[video_loop_index] + flux_prompt = "" + wan_prompt = scene.get("wan_prompt", "") + + # Matemáticas de frames exactos + duration = scene.get("duration_seconds", 5.0) + chunk_frames = math.ceil(duration * fps) + + # Ruta de donde WanVideo debe LEER la imagen inicial + image_path = os.path.join(keyframes_dir, f"scene_{scene.get('scene_id', video_loop_index)}.png") + + _log(f" -> 🎥 FASE 2 (Renderizado): Procesando Escena {scene.get('scene_id', video_loop_index)}") + _log(f" -> ⏱️ Duración de vídeo: {duration}s ({chunk_frames} frames a {fps} FPS)") + _log(f" -> 📥 Cargando Keyframe desde: {image_path}") + + _log(f"{'='*50}\n") + + return (is_flux_phase, flux_prompt, wan_prompt, chunk_frames, image_path) diff --git a/run_tests.sh b/run_tests.sh deleted file mode 100755 index e74ad2a..0000000 --- a/run_tests.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/bin/bash -set -e - -echo "Running tests in isolated processes..." -cd tests -export PYTHONPATH=.. - -pytest_cmd="/home/jules/.local/share/pipx/venvs/pytest/bin/python -m pytest -v" - -# Ignoramos test_security_fix.py temporalmente porque falla por mock the torch en otros archivos -$pytest_cmd test_install_fix.py -$pytest_cmd test_loop_calculator.py -$pytest_cmd test_ssrf_fix.py -$pytest_cmd test_switch_node.py -$pytest_cmd test_tools.py -$pytest_cmd test_video_audio.py -$pytest_cmd test_vram_node.py - -echo "All valid tests passed!" diff --git a/test_loop_syntax.py b/test_loop_syntax.py deleted file mode 100644 index 9b2b59a..0000000 --- a/test_loop_syntax.py +++ /dev/null @@ -1,24 +0,0 @@ -import sys -import unittest.mock - -# Mocking missing dependencies for syntax check -sys.modules['torch'] = unittest.mock.MagicMock() -sys.modules['comfy'] = unittest.mock.MagicMock() -sys.modules['comfy.model_management'] = unittest.mock.MagicMock() -sys.modules['folder_paths'] = unittest.mock.MagicMock() - -# Mock internal import to prevent syntax error on register_node -import types -dummy_module = types.ModuleType("dummy") -dummy_module.register_node = lambda c: c -sys.modules[__package__ or '.'] = dummy_module - -try: - with open('loop.py') as f: - code = f.read() - code = code.replace('from . import register_node', 'def register_node(c): return c') - exec(code) - print("Syntax is OK") -except Exception as e: - import traceback - traceback.print_exc() diff --git a/tests/sandbox/input/valid.csv b/tests/sandbox/input/valid.csv deleted file mode 100644 index faecd22..0000000 --- a/tests/sandbox/input/valid.csv +++ /dev/null @@ -1,2 +0,0 @@ -col1,col2 -val1,val2 \ No newline at end of file diff --git a/tests/sandbox/secret.txt b/tests/sandbox/secret.txt deleted file mode 100644 index 3146c36..0000000 --- a/tests/sandbox/secret.txt +++ /dev/null @@ -1,2 +0,0 @@ -secret_key,secret_value -super,secret \ No newline at end of file diff --git a/tests/test_install_fix.py b/tests/test_install_fix.py deleted file mode 100644 index 2bec41c..0000000 --- a/tests/test_install_fix.py +++ /dev/null @@ -1,34 +0,0 @@ -import unittest -from unittest.mock import patch, MagicMock -import sys -import os - -# Mock sys.platform to linux for testing check_and_install_ffmpeg -@patch("sys.platform", "linux") -@patch("shutil.which", return_value=None) -class TestInstall(unittest.TestCase): - - @patch("subprocess.check_call") - def test_check_and_install_ffmpeg_linux(self, mock_check_call, mock_which): - from install import check_and_install_ffmpeg - - check_and_install_ffmpeg() - - # Verify that subprocess.check_call was called twice with the correct lists - self.assertEqual(mock_check_call.call_count, 2) - - expected_calls = [ - unittest.mock.call(["sudo", "apt-get", "update"]), - unittest.mock.call(["sudo", "DEBIAN_FRONTEND=noninteractive", "apt-get", "install", "-y", "ffmpeg"]) - ] - mock_check_call.assert_has_calls(expected_calls) - - # Verify shell=True was NOT used - for call in mock_check_call.call_args_list: - args, kwargs = call - if "shell" in kwargs: - self.assertFalse(kwargs["shell"], "shell should be False if explicitly provided") - # If not in kwargs, it defaults to False for subprocess.check_call - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_loop_calculator.py b/tests/test_loop_calculator.py deleted file mode 100644 index 6b95620..0000000 --- a/tests/test_loop_calculator.py +++ /dev/null @@ -1,170 +0,0 @@ -import sys -import types -import unittest -from unittest.mock import MagicMock -import os - -# Mock required modules (as per memory) -mock_torch = types.ModuleType("torch") -mock_torch.Tensor = type("Tensor", (), {}) -mock_torch.nn = type("nn", (), {"functional": type("functional", (), {}), "Module": type("Module", (), {})}) -mock_torch.int64 = MagicMock() -mock_torch.float32 = MagicMock() -mock_torch.from_numpy = MagicMock() -mock_torch.int32 = MagicMock() -mock_torch.bfloat16 = MagicMock() -mock_torch.float16 = MagicMock() -mock_torch.cat = MagicMock() -mock_torch.zeros = MagicMock() -mock_torch.ones = MagicMock() -mock_torch.tensor = MagicMock() -mock_torch.int16 = MagicMock() -mock_torch.int8 = MagicMock() -mock_torch.float64 = MagicMock() -mock_torch.save = MagicMock() -mock_torch.load = MagicMock() -mock_torch.device = MagicMock() -mock_torch.uint8 = MagicMock() -mock_torch.Size = MagicMock() -mock_torch.autocast = MagicMock() -mock_torch.bool = MagicMock() -mock_torch.no_grad = MagicMock() -mock_torch.complex64 = MagicMock() -mock_torch.complex128 = MagicMock() -mock_torch.qint8 = MagicMock() -mock_torch.quint8 = MagicMock() -mock_torch.__version__ = "2.0.0" -mock_torch.dtype = type("dtype", (), {}) -mock_torch.empty = MagicMock() -mock_torch.randn = MagicMock() -mock_torch.dtype = type("dtype", (), {}) -mock_torch.empty = MagicMock() -mock_torch.randn = MagicMock() -mock_torch.cuda = MagicMock() -mock_torch.cuda.is_available.return_value = False -sys.modules["torch"] = mock_torch - -sys.modules["comfy"] = types.ModuleType("comfy") -sys.modules["comfy.model_management"] = types.ModuleType("comfy.model_management") -sys.modules["comfy.utils"] = types.ModuleType("comfy.utils") -sys.modules["torchaudio"] = types.ModuleType("torchaudio") -sys.modules["folder_paths"] = types.ModuleType("folder_paths") - - - -# Mock the loop module since it will be imported in video.py -mock_loop = types.ModuleType("loop") -mock_loop.global_step_by_chunk = False -mock_loop.global_source_frame_count = 0 -mock_loop.global_select_every_nth = 1 -mock_loop.global_is_final_chunk = False -mock_loop.global_accumulated_frames = 0 -sys.modules["loop"] = mock_loop - -import sys -sys.modules["tqdm"] = types.ModuleType("tqdm") -sys.modules["tqdm"].tqdm = lambda x, **kwargs: x -sys.modules["cv2"] = types.ModuleType("cv2") - - -def get_node_class(file_name, class_name): - # Read the file - path = os.path.join(os.path.dirname(__file__), "..", file_name) - with open(path, 'r') as f: - code = f.read() - - # Mock 'from . import register_node' - code = code.replace("from . import register_node", "def register_node(c, *args, **kwargs): return c") - # In video.py, there are more relative imports - code = code.replace("from . import loop", "import loop") - - # We create a new module - module = types.ModuleType(file_name.replace(".py", "")) - - # We must patch sys.modules inside so it can import loop - sys.modules[module.__name__] = module - - try: - exec(code, module.__dict__) - except Exception as e: - print(f"Failed to execute {file_name}: {e}") - pass - - return getattr(module, class_name, None) - - -class TestAutoLoopCalculators(unittest.TestCase): - - def setUp(self): - LTXClass = get_node_class("loop.py", "AutoLoopCalculatorLTX") - self.calc_ltx = LTXClass() if LTXClass else None - - BaseClass = get_node_class("video.py", "AutoLoopCalculator") - self.calc_base = BaseClass() if BaseClass else None - - WanClass = get_node_class("video.py", "AutoLoopCalculatorWan") - self.calc_wan = WanClass() if WanClass else None - - # LTX Calculator - def test_ltx_guard_zero_source(self): - res = self.calc_ltx.calculate(0, 100, 1, 0) - self.assertEqual(res, (1, 0, 1)) - - def test_ltx_guard_negative_source(self): - res = self.calc_ltx.calculate(-5, 100, 1, 0) - self.assertEqual(res, (1, 0, 1)) - - def test_ltx_guard_zero_target(self): - res = self.calc_ltx.calculate(100, 0, 1, 0) - self.assertEqual(res, (1, 0, 1)) - - def test_ltx_valid_input(self): - res = self.calc_ltx.calculate(80, 81, 1, 0) - self.assertEqual(res[0], 81) - self.assertEqual(res[1], 0) - self.assertEqual(res[2], 1) - - # Base Calculator - def test_base_guard_zero_source(self): - res = self.calc_base.calculate(0, 100, 1, 0) - self.assertEqual(res[0], 1) - self.assertEqual(res[1], 0) - self.assertEqual(res[2], 1) - self.assertIn("Invalid frame count", res[3]) - - def test_base_guard_zero_target(self): - res = self.calc_base.calculate(100, 0, 1, 0) - self.assertEqual(res[0], 1) - self.assertEqual(res[1], 0) - self.assertEqual(res[2], 1) - self.assertIn("Invalid frame count", res[3]) - - def test_base_valid_input(self): - res = self.calc_base.calculate(100, 50, 1, 0) - self.assertEqual(res[0], 50) - self.assertEqual(res[1], 0) - self.assertEqual(res[2], 1) - - # Wan Calculator - def test_wan_guard_zero_source(self): - res = self.calc_wan.calculate(0, 100, 1, 0) - self.assertEqual(res[0], 1) - self.assertEqual(res[1], 0) - self.assertEqual(res[2], 1) - self.assertIn("Invalid frame count", res[3]) - - def test_wan_guard_zero_target(self): - res = self.calc_wan.calculate(100, 0, 1, 0) - self.assertEqual(res[0], 1) - self.assertEqual(res[1], 0) - self.assertEqual(res[2], 1) - self.assertIn("Invalid frame count", res[3]) - - def test_wan_valid_input(self): - res = self.calc_wan.calculate(80, 81, 1, 0) - self.assertEqual(res[0], 81) - self.assertEqual(res[1], 0) - self.assertEqual(res[2], 1) - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_security_fix.py b/tests/test_security_fix.py deleted file mode 100644 index 92e3436..0000000 --- a/tests/test_security_fix.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch -import os -import tempfile -import unittest -from safetensors.torch import save_file, load_file - -class TestTorchLoadSecurity(unittest.TestCase): - def test_safetensors_serialization(self): - # Create a dummy tensor and save it - data = { - "waveform": torch.randn(1, 44100), - "sample_rate": torch.tensor(44100, dtype=torch.int32) - } - - with tempfile.TemporaryDirectory() as tmpdir: - file_path = os.path.join(tmpdir, "test.safetensors") - save_file(data, file_path) - - # Safetensors completely eliminates pickle - loaded_data = load_file(file_path) - - self.assertEqual(loaded_data["sample_rate"].item(), 44100) - self.assertTrue(torch.equal(loaded_data["waveform"], data["waveform"])) - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_ssrf_fix.py b/tests/test_ssrf_fix.py deleted file mode 100644 index 2f7067a..0000000 --- a/tests/test_ssrf_fix.py +++ /dev/null @@ -1,46 +0,0 @@ -import sys -import unittest -import types - -class TestSSRF(unittest.TestCase): - def test_ssrf_fix(self): - # We want to load loop.py and simulate the port logic. - import os - import urllib.request - from unittest.mock import patch, MagicMock - - # Load loop.py dynamically - REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - loop_path = os.path.join(REPO_ROOT, "loop.py") - with open(loop_path, 'r') as f: - code = f.read() - - # Mock register_node by replacing the import - code = code.replace("from . import register_node", "def register_node(c, *args, **kwargs): return c") - - # Mock comfy modules - sys.modules["comfy"] = MagicMock() - sys.modules["comfy.model_management"] = MagicMock() - sys.modules["torch"] = MagicMock() - - module = types.ModuleType("loop") - exec(code, module.__dict__) - - SequentialLoopStart = module.SequentialLoopStart() - - # Scenario 1: Malicious port input, should fallback to sys.argv or default - sys.argv = ["main.py", "--port", "8188"] - with patch('urllib.request.urlopen') as mock_urlopen: - mock_urlopen.return_value = MagicMock() - idx = SequentialLoopStart.get_index(reset_loop=True, loop_idx=0, port=9999) - - # The safe port should be 8188 since sys.argv overrides the malicious port - self.assertEqual(module.global_server_port, 8188) - - # Scenario 2: sys.argv contains a different valid port - sys.argv = ["main.py", "--port", "8189"] - idx = SequentialLoopStart.get_index(reset_loop=True, loop_idx=0, port=9999) - self.assertEqual(module.global_server_port, 8189) - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_switch_node.py b/tests/test_switch_node.py deleted file mode 100644 index 5b5cd03..0000000 --- a/tests/test_switch_node.py +++ /dev/null @@ -1,44 +0,0 @@ -import sys -import os -import pytest - -# Add the parent directory to sys.path so we can import switch_node directly -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) - -from switch_node import MasterSwitch - -@pytest.fixture -def master_switch(): - return MasterSwitch() - -def test_master_switch_check_lazy_status(master_switch): - # Test True condition - assert master_switch.check_lazy_status(is_final_cycle=True) == ["on_true"] - - # Test False condition - assert master_switch.check_lazy_status(is_final_cycle=False) == ["on_false"] - -def test_master_switch_route_true(master_switch): - # Test routing to true branch - result = master_switch.route(is_final_cycle=True, on_true="true_value", on_false="false_value") - - # Verify result - check first element to avoid tuple length mismatch assumptions - assert result[0] == "true_value" - -def test_master_switch_route_false(master_switch): - # Test routing to false branch - result = master_switch.route(is_final_cycle=False, on_true="true_value", on_false="false_value") - - # Verify result - assert result[0] == "false_value" - -def test_master_switch_route_none_values(master_switch): - # Test with None values - result_true = master_switch.route(is_final_cycle=True, on_true=None, on_false="false_value") - assert result_true[0] is None - - result_false = master_switch.route(is_final_cycle=False, on_true="true_value", on_false=None) - assert result_false[0] is None - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_tools.py b/tests/test_tools.py deleted file mode 100644 index 1418f37..0000000 --- a/tests/test_tools.py +++ /dev/null @@ -1,50 +0,0 @@ -import sys -import types -import pytest -from unittest.mock import MagicMock -import os - -# Mock dependencies -mock_torch = types.ModuleType("torch") -class MockTensor: - pass -mock_torch.Tensor = MockTensor -mock_torch.Tensor.__module__ = "torch" -sys.modules["torch"] = mock_torch - -def get_VAESafeFramePadder(): - tools_path = os.path.join(os.path.dirname(__file__), "..", "tools.py") - with open(tools_path, 'r') as f: - code = f.read() - - # Mock 'from . import register_node' and other potential relative imports - code = code.replace("from . import register_node", "def register_node(c, *args, **kwargs): return c") - code = code.replace("from .switch_node import any_type", "any_type = '*'") - - tools_module = types.ModuleType("tools") - # Provide torch to the module's namespace - tools_module.__dict__["torch"] = mock_torch - - try: - exec(code, tools_module.__dict__) - except Exception as e: - print(f"Error executing tools.py code: {e}") - raise e - - return tools_module.VAESafeFramePadder - -@pytest.fixture -def VAESafeFramePadder(): - return get_VAESafeFramePadder() - -def test_pad_non_tensor_input_raises_value_error(VAESafeFramePadder): - padder = VAESafeFramePadder() - - with pytest.raises(ValueError, match=r"\[VAESafeFramePadder\] Input is not a valid tensor\."): - padder.pad(None, "WanVideo") - - with pytest.raises(ValueError, match=r"\[VAESafeFramePadder\] Input is not a valid tensor\."): - padder.pad("not a tensor", "WanVideo") - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/test_video_audio.py b/tests/test_video_audio.py deleted file mode 100644 index 1808f11..0000000 --- a/tests/test_video_audio.py +++ /dev/null @@ -1,162 +0,0 @@ -import sys -import types -import pytest -from unittest.mock import MagicMock, PropertyMock, patch -import os - -# Create persistent mocks that we will reset in each test -m_torch = MagicMock() -m_torch.__version__ = "2.1.0" -m_torchaudio = MagicMock() -m_transforms = MagicMock() - -# Mock dependencies before importing anything that might depend on them -sys.modules["torch"] = m_torch -sys.modules["torchaudio"] = m_torchaudio -sys.modules["torchaudio.transforms"] = m_transforms - -sys.modules["folder_paths"] = MagicMock() -sys.modules["nodes"] = MagicMock() -sys.modules["comfy"] = MagicMock() -sys.modules["comfy.utils"] = MagicMock() -sys.modules["comfy.model_management"] = MagicMock() -sys.modules["tqdm"] = MagicMock() -sys.modules["PIL"] = MagicMock() -sys.modules["cv2"] = MagicMock() -sys.modules["numpy"] = MagicMock() - -def get_extract_and_standardize_audio(): - # Read the file - video_path = os.path.join(os.path.dirname(__file__), "..", "video.py") - with open(video_path, 'r', encoding='utf-8') as f: - code = f.read() - - # Mock 'from . import register_node' by replacing it - code = code.replace("from . import register_node", "def register_node(c, *args, **kwargs): return c") - - # ENSURE import torchaudio.transforms as T uses our mock - # We can do this by replacing the string in the code - code = code.replace("import torchaudio.transforms as T", "T = m_transforms") - - video_module = types.ModuleType("video") - # We need to set some things in the module dict that might be needed - video_module.__dict__.update({ - "__file__": video_path, - "torch": m_torch, - "torchaudio": m_torchaudio, - "m_transforms": m_transforms, - "os": os, - "sys": sys, - "logging": MagicMock(), - "ThreadPoolExecutor": MagicMock(), - "tqdm": MagicMock(), - "comfy": sys.modules["comfy"], - "folder_paths": sys.modules["folder_paths"], - "nodes": sys.modules["nodes"], - "subprocess": MagicMock(), - "math": MagicMock(), - "time": MagicMock(), - "uuid": MagicMock(), - "redirect_stdout": MagicMock(), - "redirect_stderr": MagicMock(), - }) - - exec(code, video_module.__dict__) - - return video_module.extract_and_standardize_audio - -# We will load the function ONCE at the module level to avoid re-execing video.py -# which might have side effects on sys.modules or other things. -_extract_and_standardize_audio = get_extract_and_standardize_audio() - -@pytest.fixture -def func(): - # Reset mocks before each test - m_torch.reset_mock() - m_torchaudio.reset_mock() - m_transforms.reset_mock() - - # Re-ensure torch.mean exists after reset as it's used in the function - m_torch.mean = MagicMock() - - # Return the already loaded function - return _extract_and_standardize_audio - -def test_extract_and_standardize_audio_exception(func): - # Mock torchaudio.load to raise an Exception - m_torchaudio.load.side_effect = Exception("Test exception") - - with pytest.raises(Exception) as excinfo: - func("dummy_path") - - assert str(excinfo.value) == "Test exception" - m_torchaudio.load.assert_called_once_with("dummy_path") - -def test_extract_and_standardize_audio_happy_path_no_resample_stereo(func): - m_torchaudio.load.side_effect = None - # Mock waveform: 2 channels, some length - mock_waveform = MagicMock() - # Use PropertyMock for shape to avoid it being a MagicMock that causes issues in comparison - type(mock_waveform).shape = PropertyMock(return_value=(2, 1000)) - m_torchaudio.load.return_value = (mock_waveform, int(44100)) - - result = func("dummy_path", target_sr=44100) - - assert result["sample_rate"] == 44100 - mock_waveform.unsqueeze.assert_called_once_with(0) - assert result["waveform"] == mock_waveform.unsqueeze.return_value - -def test_extract_and_standardize_audio_resample(func): - m_torchaudio.load.side_effect = None - mock_waveform = MagicMock() - type(mock_waveform).shape = PropertyMock(return_value=(2, 1000)) - m_torchaudio.load.return_value = (mock_waveform, int(48000)) - - mock_resampled_waveform = MagicMock() - type(mock_resampled_waveform).shape = PropertyMock(return_value=(2, 918)) - - mock_resampler = MagicMock() - mock_resampler.return_value = mock_resampled_waveform - m_transforms.Resample.return_value = mock_resampler - - result = func("dummy_path", target_sr=44100) - - m_transforms.Resample.assert_called_once_with(orig_freq=48000, new_freq=44100) - mock_resampler.assert_called_once_with(mock_waveform) - assert result["sample_rate"] == 44100 - assert result["waveform"] == mock_resampled_waveform.unsqueeze.return_value - -def test_extract_and_standardize_audio_mono_to_stereo(func): - m_torchaudio.load.side_effect = None - mock_waveform = MagicMock() - type(mock_waveform).shape = PropertyMock(return_value=(1, 1000)) - m_torchaudio.load.return_value = (mock_waveform, int(44100)) - - mock_stereo_waveform = MagicMock() - type(mock_stereo_waveform).shape = PropertyMock(return_value=(2, 1000)) - mock_waveform.repeat.return_value = mock_stereo_waveform - - result = func("dummy_path", target_sr=44100) - - mock_waveform.repeat.assert_called_once_with(2, 1) - assert result["waveform"] == mock_stereo_waveform.unsqueeze.return_value - -def test_extract_and_standardize_audio_multichannel_to_stereo(func): - m_torchaudio.load.side_effect = None - mock_waveform = MagicMock() - type(mock_waveform).shape = PropertyMock(return_value=(6, 1000)) - m_torchaudio.load.return_value = (mock_waveform, int(44100)) - - mock_mono = MagicMock() - type(mock_mono).shape = PropertyMock(return_value=(1, 1000)) - m_torch.mean.return_value = mock_mono - - mock_stereo = MagicMock() - type(mock_stereo).shape = PropertyMock(return_value=(2, 1000)) - mock_mono.repeat.return_value = mock_stereo - - result = func("dummy_path", target_sr=44100) - - m_torch.mean.assert_called_once_with(mock_waveform, dim=0, keepdim=True) - mock_mono.repeat.assert_called_once_with(2, 1) - assert result["waveform"] == mock_stereo.unsqueeze.return_value diff --git a/tests/test_vram_node.py b/tests/test_vram_node.py deleted file mode 100644 index 35f4c78..0000000 --- a/tests/test_vram_node.py +++ /dev/null @@ -1,55 +0,0 @@ -import sys -import types -import pytest -from unittest.mock import MagicMock -import os - -# Mock dependencies before importing anything that might depend on them -mock_torch = types.ModuleType("torch") -mock_torch.cuda = MagicMock() -mock_torch.cuda.is_available.return_value = False -sys.modules["torch"] = mock_torch - -sys.modules["comfy"] = types.ModuleType("comfy") -sys.modules["comfy.model_management"] = types.ModuleType("comfy.model_management") - -def get_format_bytes(): - # Read the file - vram_node_path = os.path.join(os.path.dirname(__file__), "..", "vram_node.py") - with open(vram_node_path, 'r') as f: - code = f.read() - - # Mock 'from . import register_node' by replacing it - code = code.replace("from . import register_node", "def register_node(c, *args, **kwargs): return c") - - vram_module = types.ModuleType("vram_node") - exec(code, vram_module.__dict__) - - return vram_module.format_bytes - -@pytest.fixture -def format_bytes(): - return get_format_bytes() - -def test_format_bytes_zero(format_bytes): - assert format_bytes(0) == "0.00 GB" - -def test_format_bytes_gigabyte(format_bytes): - assert format_bytes(1024**3) == "1.00 GB" - -def test_format_bytes_fractional(format_bytes): - assert format_bytes(1024**3 * 0.5) == "0.50 GB" - assert format_bytes(1024**3 * 2.5) == "2.50 GB" - -def test_format_bytes_small_values(format_bytes): - # 1 MB is 1/1024 GB ≈ 0.000976 GB -> 0.00 GB - assert format_bytes(1024**2) == "0.00 GB" - # 10 MB is 10/1024 GB ≈ 0.009765 GB -> 0.01 GB (due to rounding .2f) - assert format_bytes(1024**2 * 10) == "0.01 GB" - -def test_format_bytes_large_values(format_bytes): - assert format_bytes(1024**3 * 1024) == "1024.00 GB" - assert format_bytes(1024**3 * 123.456) == "123.46 GB" - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/tests/verify_fix.py b/tests/verify_fix.py deleted file mode 100644 index f66fe29..0000000 --- a/tests/verify_fix.py +++ /dev/null @@ -1,118 +0,0 @@ -import os -import sys -import csv - -# Setup paths -TEST_DIR = os.path.dirname(os.path.abspath(__file__)) -REPO_ROOT = os.path.dirname(TEST_DIR) -SANDBOX_DIR = os.path.join(TEST_DIR, "sandbox") -INPUT_DIR = os.path.join(SANDBOX_DIR, "input") -SECRET_FILE = os.path.join(SANDBOX_DIR, "secret.txt") - -# Ensure sandbox exists -if not os.path.exists(INPUT_DIR): - os.makedirs(INPUT_DIR) - -# Create secret file -with open(SECRET_FILE, 'w') as f: - f.write("secret_key,secret_value\nsuper,secret") - -# Create valid csv -with open(os.path.join(INPUT_DIR, "valid.csv"), 'w') as f: - f.write("col1,col2\nval1,val2") - -# Mock folder_paths -import types -folder_paths = types.ModuleType("folder_paths") -folder_paths.get_input_directory = lambda: INPUT_DIR -folder_paths.get_annotated_filepath = lambda x: x # Mock implementation -sys.modules["folder_paths"] = folder_paths - -# Load batch.py dynamically -def load_batch_module(): - batch_path = os.path.join(REPO_ROOT, "batch.py") - with open(batch_path, 'r') as f: - code = f.read() - - # Mock register_node by replacing the import - code = code.replace("from . import register_node", "def register_node(c, *args, **kwargs): return c") - - module = types.ModuleType("batch") - exec(code, module.__dict__) - return module - -batch_module = load_batch_module() -LoadCSV = batch_module.LoadCSV - -def verify_fix(): - print("--- Verifying Fix ---") - node = LoadCSV() - - # Test 1: Exploit attempt (should fail) - print("Test 1: Exploit attempt (should raise ValueError)") - target_path = SECRET_FILE - try: - # Pass path relative to input, attempting traversal if we were naive, - # but here we can pass absolute path to check absolute path blocking. - # Or relative path "../secret.txt" - - # Scenario A: Path traversal relative string - # node.go("../secret.txt", ...) - # Scenario B: Absolute path - - # Testing Scenario B (absolute path) - node.go(target_path, "comma", '"', -1) - print("FAIL: Exploit succeeded (No exception raised)") - return False - except ValueError as e: - if "Access denied" in str(e): - print(f"SUCCESS: Blocked absolute path access: {e}") - else: - print(f"FAIL: Raised wrong ValueError: {e}") - return False - except Exception as e: - print(f"FAIL: Raised wrong exception: {type(e)} {e}") - return False - - # Test 2: Valid access (should succeed) - print("\nTest 2: Valid access (should succeed)") - try: - # Can pass just filename or relative path - result = node.go("valid.csv", "comma", '"', -1) - batch_data = result["result"][0] - if batch_data[0]['col1'] == 'val1': - print("SUCCESS: Valid CSV read correctly") - else: - print(f"FAIL: content mismatch: {batch_data}") - return False - except Exception as e: - print(f"FAIL: Valid access raised exception: {e}") - return False - - # Test 3: Relative path traversal attempt (should fail) - print("\nTest 3: Relative path traversal attempt (should raise ValueError)") - try: - node.go("../secret.txt", "comma", '"', -1) - print("FAIL: Exploit succeeded (No exception raised)") - return False - except ValueError as e: - if "Access denied" in str(e): - print(f"SUCCESS: Blocked relative path traversal: {e}") - else: - print(f"FAIL: Raised wrong ValueError: {e}") - return False - except Exception as e: - # Depending on how os.path.join works, "input/../secret.txt" -> "sandbox/secret.txt" - # commonpath should catch it. - print(f"FAIL: Raised wrong exception: {type(e)} {e}") - return False - - return True - -if __name__ == "__main__": - if verify_fix(): - print("\nVerification Result: FIX VERIFIED") - sys.exit(0) - else: - print("\nVerification Result: FIX FAILED") - sys.exit(1)