403 lines
14 KiB
Python
403 lines
14 KiB
Python
import pytest
|
|
import os
|
|
import tempfile
|
|
import shutil
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from unittest.mock import patch, Mock, MagicMock
|
|
|
|
# Import the functions we're testing
|
|
try:
|
|
from nodes import list_files_from_directory, load_image_video_from_path, select_indexes_from_any
|
|
except ImportError:
|
|
# Fallback for when running tests directly
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
|
from nodes import list_files_from_directory, load_image_video_from_path, select_indexes_from_any
|
|
|
|
|
|
class TestListFilesFromDirectory:
|
|
"""Test cases for list_files_from_directory function."""
|
|
|
|
def test_basic_file_listing(self, test_files_directory):
|
|
result = list_files_from_directory(
|
|
directory=test_files_directory,
|
|
file_list_cap=0,
|
|
file_start_index=0,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg,png',
|
|
sort_method='None'
|
|
)
|
|
filenames = result[0]
|
|
|
|
# Should find image files but not video or txt files
|
|
assert len(filenames) >= 2
|
|
assert any('image1.jpg' in f for f in filenames)
|
|
assert any('image2.png' in f for f in filenames)
|
|
assert not any('video.mp4' in f for f in filenames)
|
|
assert not any('document.txt' in f for f in filenames)
|
|
|
|
def test_include_subfolders(self, test_files_directory):
|
|
result = list_files_from_directory(
|
|
directory=test_files_directory,
|
|
file_list_cap=0,
|
|
file_start_index=0,
|
|
include_subfolders=True,
|
|
valid_extensions='jpg,png',
|
|
sort_method='None'
|
|
)
|
|
filenames = result[0]
|
|
|
|
# Should find files in subdirectories
|
|
assert any('nested.jpg' in f for f in filenames)
|
|
|
|
def test_file_list_cap(self, test_files_directory):
|
|
result = list_files_from_directory(
|
|
directory=test_files_directory,
|
|
file_list_cap=1,
|
|
file_start_index=0,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg,png',
|
|
sort_method='None'
|
|
)
|
|
filenames = result[0]
|
|
|
|
# Should limit to 1 file
|
|
assert len(filenames) == 1
|
|
|
|
def test_file_start_index(self, test_files_directory):
|
|
# First get all files
|
|
all_result = list_files_from_directory(
|
|
directory=test_files_directory,
|
|
file_list_cap=0,
|
|
file_start_index=0,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg,png',
|
|
sort_method='Alphabetical (ASC)'
|
|
)
|
|
all_files = all_result[0]
|
|
|
|
if len(all_files) > 1:
|
|
# Now get files starting from index 1
|
|
result = list_files_from_directory(
|
|
directory=test_files_directory,
|
|
file_list_cap=0,
|
|
file_start_index=1,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg,png',
|
|
sort_method='Alphabetical (ASC)'
|
|
)
|
|
subset_files = result[0]
|
|
|
|
# Should be missing the first file
|
|
assert len(subset_files) == len(all_files) - 1
|
|
assert subset_files[0] == all_files[1]
|
|
|
|
def test_valid_extensions_parsing(self, test_files_directory):
|
|
# Test various extension formats
|
|
result = list_files_from_directory(
|
|
directory=test_files_directory,
|
|
file_list_cap=0,
|
|
file_start_index=0,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg, png , mp4', # With spaces
|
|
sort_method='None'
|
|
)
|
|
filenames = result[0]
|
|
|
|
# Should find jpg, png, and mp4 files
|
|
extensions_found = set()
|
|
for f in filenames:
|
|
ext = f.split('.')[-1].lower()
|
|
extensions_found.add(ext)
|
|
|
|
expected = {'jpg', 'png', 'mp4'} & extensions_found
|
|
assert len(expected) > 0
|
|
|
|
def test_nonexistent_directory(self):
|
|
with pytest.raises(FileNotFoundError, match="Directory.*cannot be found"):
|
|
list_files_from_directory(
|
|
directory='/nonexistent/directory',
|
|
file_list_cap=0,
|
|
file_start_index=0,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg,png',
|
|
sort_method='None'
|
|
)
|
|
|
|
def test_empty_directory_no_matching_files(self, temp_dir):
|
|
# Create directory with only non-matching files
|
|
with open(os.path.join(temp_dir, 'test.txt'), 'w') as f:
|
|
f.write('content')
|
|
|
|
with pytest.raises(FileNotFoundError, match="No files in directory"):
|
|
list_files_from_directory(
|
|
directory=temp_dir,
|
|
file_list_cap=0,
|
|
file_start_index=0,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg,png',
|
|
sort_method='None'
|
|
)
|
|
|
|
def test_sorting_integration(self, test_files_directory):
|
|
with patch('nodes.sort_by') as mock_sort:
|
|
mock_sort.return_value = ['sorted', 'files']
|
|
|
|
list_files_from_directory(
|
|
directory=test_files_directory,
|
|
file_list_cap=0,
|
|
file_start_index=0,
|
|
include_subfolders=False,
|
|
valid_extensions='jpg,png',
|
|
sort_method='Alphabetical (ASC)'
|
|
)
|
|
|
|
# Should call sort_by with the method
|
|
mock_sort.assert_called_once()
|
|
args = mock_sort.call_args[0]
|
|
kwargs = mock_sort.call_args[1] if mock_sort.call_args[1] else {}
|
|
|
|
assert kwargs.get('method') == 'Alphabetical (ASC)' or args[2] == 'Alphabetical (ASC)'
|
|
|
|
|
|
class TestLoadImageVideoFromPath:
|
|
"""Test cases for load_image_video_from_path function."""
|
|
|
|
def create_test_image(self, path, mode='RGB', size=(100, 100), color='red'):
|
|
"""Helper to create test images."""
|
|
img = Image.new(mode, size, color)
|
|
img.save(path)
|
|
return img
|
|
|
|
def test_load_simple_rgb_image(self, temp_dir):
|
|
image_path = os.path.join(temp_dir, 'test.png')
|
|
self.create_test_image(image_path, mode='RGB', size=(50, 50), color=(255, 0, 0))
|
|
|
|
with patch('nodes.node_helpers.pillow') as mock_pillow:
|
|
mock_pillow.side_effect = lambda func, *args: func(*args)
|
|
|
|
result = load_image_video_from_path(path=image_path, RGBA=False)
|
|
output_image, output_mask, audio, frames, fps = result
|
|
|
|
assert isinstance(output_image, torch.Tensor)
|
|
assert output_image.shape[-1] == 3 # RGB channels
|
|
assert frames == 1
|
|
assert fps == 0.0
|
|
assert audio is None
|
|
|
|
def test_load_rgba_image_with_alpha(self, temp_dir):
|
|
image_path = os.path.join(temp_dir, 'test.png')
|
|
self.create_test_image(image_path, mode='RGBA', size=(50, 50), color=(255, 0, 0, 128))
|
|
|
|
with patch('nodes.node_helpers.pillow') as mock_pillow:
|
|
mock_pillow.side_effect = lambda func, *args: func(*args)
|
|
|
|
result = load_image_video_from_path(path=image_path, RGBA=True)
|
|
output_image, output_mask, audio, frames, fps = result
|
|
|
|
assert isinstance(output_image, torch.Tensor)
|
|
assert output_image.shape[-1] == 4 # RGBA channels
|
|
assert isinstance(output_mask, torch.Tensor)
|
|
|
|
def test_load_video_file(self, temp_dir):
|
|
video_path = os.path.join(temp_dir, 'test.mp4')
|
|
|
|
# Mock VideoFromFile
|
|
mock_video = Mock()
|
|
mock_components = Mock()
|
|
mock_components.images = torch.rand(10, 100, 100, 3)
|
|
mock_components.audio = torch.rand(1000, 2)
|
|
mock_components.frame_rate = 30.0
|
|
mock_video.get_components.return_value = mock_components
|
|
|
|
with patch('nodes.VideoFromFile', return_value=mock_video):
|
|
result = load_image_video_from_path(path=video_path)
|
|
output_image, output_mask, audio, frames, fps = result
|
|
|
|
assert isinstance(output_image, torch.Tensor)
|
|
assert output_image.shape[0] == 10 # 10 frames
|
|
assert fps == 30.0
|
|
assert audio is not None
|
|
|
|
def test_resize_functionality(self, temp_dir):
|
|
image_path = os.path.join(temp_dir, 'test.png')
|
|
self.create_test_image(image_path, size=(200, 100))
|
|
|
|
with patch('nodes.node_helpers.pillow') as mock_pillow, \
|
|
patch('nodes.resize_with_aspect_ratio') as mock_resize:
|
|
|
|
mock_pillow.side_effect = lambda func, *args: func(*args)
|
|
mock_resized = Image.new('RGB', (150, 150), 'blue')
|
|
mock_resize.return_value = mock_resized
|
|
|
|
load_image_video_from_path(
|
|
path=image_path,
|
|
width=150,
|
|
height=150,
|
|
keep_aspect_ratio='crop'
|
|
)
|
|
|
|
mock_resize.assert_called_once()
|
|
args = mock_resize.call_args[0]
|
|
assert args[1] == 150 # width
|
|
assert args[2] == 150 # height
|
|
assert args[3] == 'crop' # keep_aspect_ratio
|
|
|
|
def test_animated_gif_handling(self, temp_dir):
|
|
# Create a simple animated GIF
|
|
gif_path = os.path.join(temp_dir, 'test.gif')
|
|
|
|
# Create frames
|
|
frames = []
|
|
for i in range(3):
|
|
frame = Image.new('RGB', (50, 50), (i * 80, 0, 0))
|
|
frame.info['duration'] = 100 # 100ms per frame
|
|
frames.append(frame)
|
|
|
|
frames[0].save(
|
|
gif_path,
|
|
save_all=True,
|
|
append_images=frames[1:],
|
|
duration=100
|
|
)
|
|
|
|
with patch('nodes.node_helpers.pillow') as mock_pillow:
|
|
mock_pillow.side_effect = lambda func, *args: func(*args)
|
|
|
|
result = load_image_video_from_path(path=gif_path)
|
|
output_image, output_mask, audio, frames_count, fps = result
|
|
|
|
assert isinstance(output_image, torch.Tensor)
|
|
assert frames_count == 3
|
|
assert fps > 0 # Should calculate FPS from duration
|
|
|
|
def test_mime_type_detection(self, temp_dir):
|
|
# Test with unknown extension
|
|
# Create a valid PNG first, then copy to unknown extension
|
|
normal_path = os.path.join(temp_dir, 'test.png')
|
|
weird_path = os.path.join(temp_dir, 'test.unknown')
|
|
self.create_test_image(normal_path)
|
|
|
|
# Copy the valid image to unknown extension
|
|
import shutil
|
|
shutil.copy2(normal_path, weird_path)
|
|
|
|
with patch('mimetypes.guess_type', return_value=(None, None)), \
|
|
patch('nodes.node_helpers.pillow') as mock_pillow:
|
|
|
|
mock_pillow.side_effect = lambda func, *args: func(*args)
|
|
|
|
# Should still work by falling back to image processing
|
|
result = load_image_video_from_path(path=weird_path)
|
|
output_image, output_mask, audio, frames, fps = result
|
|
|
|
assert isinstance(output_image, torch.Tensor)
|
|
|
|
|
|
class TestSelectIndexesFromAny:
|
|
"""Test cases for select_indexes_from_any function."""
|
|
|
|
def test_torch_tensor_single_index(self):
|
|
tensor = torch.rand(5, 10, 10, 3)
|
|
result = select_indexes_from_any(tensor, "2")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 1
|
|
assert torch.equal(result[0], tensor[2])
|
|
|
|
def test_torch_tensor_multiple_indexes(self):
|
|
tensor = torch.rand(5, 10, 10, 3)
|
|
result = select_indexes_from_any(tensor, "0,2,4")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 3
|
|
assert torch.equal(result[0], tensor[0])
|
|
assert torch.equal(result[1], tensor[2])
|
|
assert torch.equal(result[2], tensor[4])
|
|
|
|
def test_torch_tensor_range_selection(self):
|
|
tensor = torch.rand(10, 5, 5, 3)
|
|
result = select_indexes_from_any(tensor, "2:5")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 3 # indexes 2, 3, 4
|
|
assert torch.equal(result, tensor[2:5])
|
|
|
|
def test_torch_tensor_open_range(self):
|
|
tensor = torch.rand(5, 10, 10, 3)
|
|
result = select_indexes_from_any(tensor, "2:")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 3 # indexes 2, 3, 4
|
|
assert torch.equal(result, tensor[2:])
|
|
|
|
def test_torch_tensor_negative_indexes(self):
|
|
tensor = torch.rand(5, 10, 10, 3)
|
|
result = select_indexes_from_any(tensor, "-1")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 1
|
|
assert torch.equal(result[0], tensor[-1])
|
|
|
|
def test_list_single_index(self):
|
|
test_list = ['a', 'b', 'c', 'd', 'e']
|
|
result = select_indexes_from_any(test_list, "2")
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert result[0] == 'c'
|
|
|
|
def test_list_multiple_indexes(self):
|
|
test_list = ['a', 'b', 'c', 'd', 'e']
|
|
result = select_indexes_from_any(test_list, "0,2,4")
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1 # itemgetter returns tuple in a list
|
|
assert result[0] == ('a', 'c', 'e')
|
|
|
|
def test_list_range_selection(self):
|
|
test_list = ['a', 'b', 'c', 'd', 'e']
|
|
result = select_indexes_from_any(test_list, "1:4")
|
|
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert result[0] == ('b', 'c', 'd')
|
|
|
|
def test_step_selection(self):
|
|
tensor = torch.rand(10, 5, 5, 3)
|
|
result = select_indexes_from_any(tensor, "0:6:2")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 3 # indexes 0, 2, 4
|
|
assert torch.equal(result[0], tensor[0])
|
|
assert torch.equal(result[1], tensor[2])
|
|
assert torch.equal(result[2], tensor[4])
|
|
|
|
def test_complex_selection(self):
|
|
tensor = torch.rand(10, 5, 5, 3)
|
|
result = select_indexes_from_any(tensor, "0,2:5,7")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 5 # indexes 0, 2, 3, 4, 7
|
|
|
|
def test_empty_selection_raises_error(self):
|
|
tensor = torch.rand(5, 10, 10, 3)
|
|
|
|
with pytest.raises(ValueError, match="No indexes selected"):
|
|
select_indexes_from_any(tensor, "")
|
|
|
|
def test_invalid_type_raises_error(self):
|
|
invalid_obj = "not a tensor or list"
|
|
|
|
with pytest.raises(TypeError, match="not a listable type"):
|
|
select_indexes_from_any(invalid_obj, "0")
|
|
|
|
def test_whitespace_handling(self):
|
|
tensor = torch.rand(5, 10, 10, 3)
|
|
result = select_indexes_from_any(tensor, " 0 , 2 : 4 ")
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape[0] == 3 # indexes 0, 2, 3
|