[content] tests/test_dataloader.py (python)
import unittest
import numpy as np
import tempfile
import os
from pathlib import Path
import pickle

from utils.dataloader import get_dataloader
from generate_arrayrecord_dataset import create_arrayrecord_dataset


class DataloaderReproducibilityTest(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.test_data_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)

        self.num_videos = 10
        self.episode_length = 16
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        
        # Generate ArrayRecord dataset for testing
        create_arrayrecord_dataset(
            output_dir=str(self.test_data_dir),
            num_videos=self.num_videos,
            seq_len=self.episode_length,
            height=self.image_height,
            width=self.image_width,
            channels=self.image_channels,
            records_per_file=5,  # Small files for testing
            seed=42,
            num_processes=1  # Single process for testing
        )
        
        # Get list of ArrayRecord files
        self.array_record_files = [
            str(self.test_data_dir / f) 
            for f in os.listdir(self.test_data_dir) 
            if f.endswith('.array_record')
        ]

        self.fixed_seed = 42

    def test_dataloader_yields_reproducible_batches(self):
        seq_len = 8
        batch_size = 2

        dataloader1 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches1 = [next(dataloader1) for _ in range(3)]

        dataloader2 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches2 = [next(dataloader2) for _ in range(3)]

        for i, (b1, b2) in enumerate(zip(batches1, batches2)):
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")

    def test_dataloader_batch_shape_and_dtype(self):
        """Test that dataloader yields batches with correct shape and dtype."""
        seq_len = 8
        batch_size = 2

        dataloader = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        
        batch = next(dataloader)
        
        # Check shape: (batch_size, seq_len, height, width, channels)
        expected_shape = (batch_size, seq_len, self.image_height, self.image_width, self.image_channels)
        self.assertEqual(batch.shape, expected_shape)
        
        # Check dtype: should be float32 and normalized to [0, 1]
        self.assertEqual(batch.dtype, np.float32)
        self.assertTrue(np.all(batch >= 0.0) and np.all(batch <= 1.0))

    def test_dataloader_handles_different_sequence_lengths(self):
        """Test that dataloader works with different sequence lengths."""
        batch_size = 2
        
        for seq_len in [4, 8, 12]:
            dataloader = get_dataloader(
                self.array_record_files,
                seq_len,
                batch_size,
                self.image_height,
                self.image_width,
                self.image_channels,
                seed=self.fixed_seed,
            )
            
            batch = next(dataloader)
            expected_shape = (batch_size, seq_len, self.image_height, self.image_width, self.image_channels)
            self.assertEqual(batch.shape, expected_shape)


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