47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
import torch
|
|
import unittest
|
|
from nodes.loop_image_node import LoopImageNode
|
|
|
|
class TestLoopImageNode(unittest.TestCase):
|
|
def setUp(self):
|
|
self.node = LoopImageNode()
|
|
|
|
def test_single_image(self):
|
|
"""Test with a single image"""
|
|
# Create a dummy single image (1x3x64x64)
|
|
images = torch.randn(1, 3, 64, 64)
|
|
result, = self.node.execute(images)
|
|
|
|
# Result should be the same as input for a single image
|
|
self.assertEqual(result.shape, images.shape)
|
|
self.assertTrue(torch.equal(result, images))
|
|
|
|
def test_two_images(self):
|
|
"""Test with two images"""
|
|
# Create two dummy images (2x3x64x64)
|
|
images = torch.randn(2, 3, 64, 64)
|
|
result, = self.node.execute(images)
|
|
|
|
# Result should be 3 images (original 2 + first one repeated)
|
|
self.assertEqual(result.shape[0], 3)
|
|
self.assertTrue(torch.equal(result[2], images[0]))
|
|
|
|
def test_many_images(self):
|
|
"""Test with more than 2 images"""
|
|
# Create 5 dummy images (5x3x64x64)
|
|
images = torch.randn(5, 3, 64, 64)
|
|
result, = self.node.execute(images)
|
|
|
|
# Result should be 8 images (5 original + 3 reversed middle)
|
|
self.assertEqual(result.shape[0], 8)
|
|
|
|
# Check that the reversed middle section is correct
|
|
# Original: [0,1,2,3,4]
|
|
# Middle reversed: [3,2,1]
|
|
# Final: [0,1,2,3,4,3,2,1]
|
|
self.assertTrue(torch.equal(result[5], images[3]))
|
|
self.assertTrue(torch.equal(result[6], images[2]))
|
|
self.assertTrue(torch.equal(result[7], images[1]))
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |