152 lines
5.2 KiB
Python
152 lines
5.2 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
import sys
|
|
import os
|
|
import numpy as np
|
|
|
|
# Add parent directory to path for imports
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# Mock dependencies before importing nodes.video_node
|
|
mock_torch = MagicMock()
|
|
sys.modules["torch"] = mock_torch
|
|
sys.modules["folder_paths"] = MagicMock()
|
|
sys.modules["comfy"] = MagicMock()
|
|
sys.modules["comfy.cli_args"] = MagicMock()
|
|
sys.modules["comfy.utils"] = MagicMock()
|
|
sys.modules["server"] = MagicMock()
|
|
|
|
# Mock PIL
|
|
mock_pil = MagicMock()
|
|
sys.modules["PIL"] = mock_pil
|
|
sys.modules["PIL.Image"] = mock_pil
|
|
sys.modules["PIL.PngImagePlugin"] = MagicMock()
|
|
sys.modules["cv2"] = MagicMock()
|
|
|
|
# Mock shared modules' submodules to avoid dependency issues
|
|
sys.modules["shared.workflow"] = MagicMock()
|
|
sys.modules["shared.workflow.sanitizer"] = MagicMock()
|
|
sys.modules["shared.workflow.prompt_extractor"] = MagicMock()
|
|
sys.modules["shared.workflow.workflow_builder"] = MagicMock()
|
|
|
|
sys.modules["shared.discord"] = MagicMock()
|
|
sys.modules["shared.discord.webhook_client"] = MagicMock()
|
|
sys.modules["shared.discord.message_builder"] = MagicMock()
|
|
sys.modules["shared.discord.cdn_extractor"] = MagicMock()
|
|
|
|
sys.modules["shared.github_integration"] = MagicMock()
|
|
sys.modules["shared.logging_config"] = MagicMock()
|
|
sys.modules["shared.filename_utils"] = MagicMock()
|
|
sys.modules["shared.path_utils"] = MagicMock()
|
|
|
|
# Mock shared.media siblings
|
|
sys.modules["shared.media.format_utils"] = MagicMock()
|
|
sys.modules["shared.media.video_encoder"] = MagicMock()
|
|
|
|
# Note: We do NOT mock "shared", "shared.media", or "shared.media.image_processing"
|
|
# because we want to load the real code for testing.
|
|
|
|
# Define the function logic we want to verify (simulating the generator consumer)
|
|
def consume_chunks(chunks):
|
|
pil_images = []
|
|
for chunk in chunks:
|
|
# Proposed logic for nodes/video_node.py
|
|
if len(chunk.shape) == 4:
|
|
# Batched chunk (B, H, W, C)
|
|
for i in range(chunk.shape[0]):
|
|
pil_images.append(f"image_from_batch_{i}")
|
|
else:
|
|
# Single frame chunk (H, W, C)
|
|
pil_images.append("image_from_single")
|
|
return pil_images
|
|
|
|
class TestPILBatchOptimization(unittest.TestCase):
|
|
|
|
def test_consumer_logic_mixed_chunks(self):
|
|
"""Test that the consumer logic correctly handles mixed 4D and 3D chunks."""
|
|
|
|
# 1. 4D Chunk (Batch of 2)
|
|
chunk_batch = np.zeros((2, 10, 10, 3), dtype=np.uint8)
|
|
|
|
# 2. 3D Chunk (Single frame) - simulating what happens if generator yields single frame
|
|
chunk_single = np.zeros((10, 10, 3), dtype=np.uint8)
|
|
|
|
chunks = [chunk_batch, chunk_single]
|
|
|
|
# Run consumer logic
|
|
images = consume_chunks(chunks)
|
|
|
|
# Verify results
|
|
# Should have 2 from batch + 1 from single = 3 images
|
|
self.assertEqual(len(images), 3)
|
|
self.assertEqual(images[0], "image_from_batch_0")
|
|
self.assertEqual(images[1], "image_from_batch_1")
|
|
self.assertEqual(images[2], "image_from_single")
|
|
|
|
def test_process_batched_images_integration(self):
|
|
"""
|
|
Verify that we can import and run process_batched_images with mocks,
|
|
and that it chunks correctly.
|
|
"""
|
|
# Import needs to happen after mocks are set up
|
|
from shared.media.image_processing import process_batched_images
|
|
|
|
# Setup mock tensor
|
|
# We need to make sure isinstance(t, torch.Tensor) works
|
|
|
|
tensor_len = 5
|
|
batch_size = 2
|
|
|
|
# Mock slicing
|
|
def getitem(self, idx):
|
|
# idx is a slice object
|
|
start = idx.start
|
|
stop = idx.stop
|
|
if stop > tensor_len:
|
|
stop = tensor_len
|
|
size = stop - start
|
|
return f"slice_{size}"
|
|
|
|
# Create a class with __len__ and __getitem__ defined
|
|
class MockTensor:
|
|
def __len__(self):
|
|
return tensor_len
|
|
def __getitem__(self, idx):
|
|
return getitem(self, idx)
|
|
|
|
mock_torch.Tensor = MockTensor
|
|
|
|
mock_tensor = mock_torch.Tensor()
|
|
|
|
# Mock tensor_to_numpy_uint8 to return numpy arrays of appropriate shape
|
|
# It needs to return (Size, H, W, C)
|
|
with patch('shared.media.image_processing.tensor_to_numpy_uint8') as mock_t2n:
|
|
def side_effect(slice_obj):
|
|
# parse size from string "slice_N"
|
|
size = int(slice_obj.split('_')[1])
|
|
return np.zeros((size, 10, 10, 3), dtype=np.uint8)
|
|
|
|
mock_t2n.side_effect = side_effect
|
|
|
|
# Run generator
|
|
generator = process_batched_images(mock_tensor, batch_size=batch_size)
|
|
chunks = list(generator)
|
|
|
|
# Expected:
|
|
# 5 items, batch 2
|
|
# 1. Size 2
|
|
# 2. Size 2
|
|
# 3. Size 1
|
|
|
|
self.assertEqual(len(chunks), 3)
|
|
self.assertEqual(chunks[0].shape[0], 2)
|
|
self.assertEqual(chunks[1].shape[0], 2)
|
|
self.assertEqual(chunks[2].shape[0], 1)
|
|
|
|
# Verify they are all 4D arrays (B, H, W, C)
|
|
for c in chunks:
|
|
self.assertEqual(len(c.shape), 4)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|