Spaces:

readylab
/

VizRef

Sleeping

VizRef / scripts /test /test_data_utils_comprehensive.py

Add model and inference code

18a82fb about 1 year ago

22 kB

# !/usr/bin/env python3
"""
Comprehensive test script for data_utils.py
Tests the orchestration layer that combines transforms, url_dataset, and creates DataLoaders
Run from project root: python test_data_utils_comprehensive.py
"""

import sys
import os
from pathlib import Path
import torch
import numpy as np
from PIL import Image
import json
import tempfile
import traceback
import logging
from typing import Dict, List
from unittest.mock import patch, MagicMock

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')


class TestRunner:
    """Test runner with results tracking"""

    def __init__(self):
        self.passed = 0
        self.failed = 0
        self.skipped = 0
        self.errors = []

    def run_test(self, test_name: str, test_func):
        """Run a single test with error handling"""
        print(f"\n{'=' * 60}")
        print(f"Running: {test_name}")
        print(f"{'=' * 60}")

        try:
            test_func()
            self.passed += 1
            print(f"✅ PASSED: {test_name}")
        except FileNotFoundError as e:
            self.skipped += 1
            print(f"⚠️  SKIPPED: {test_name} - {str(e)}")
        except Exception as e:
            self.failed += 1
            self.errors.append({
                'test': test_name,
                'error': str(e),
                'traceback': traceback.format_exc()
            })
            print(f"❌ FAILED: {test_name}")
            print(f"Error: {str(e)}")
            print(f"Traceback:\n{traceback.format_exc()}")

    def print_summary(self):
        """Print test summary"""
        total = self.passed + self.failed + self.skipped
        print(f"\n{'=' * 60}")
        print(f"TEST SUMMARY")
        print(f"{'=' * 60}")
        print(f"Total tests: {total}")
        print(f"✅ Passed: {self.passed}")
        print(f"❌ Failed: {self.failed}")
        print(f"⚠️  Skipped: {self.skipped}")

        if self.errors:
            print(f"\nFailed tests:")
            for error in self.errors:
                print(f"  - {error['test']}: {error['error']}")


def create_mock_model_configs():
    """Create mock model configs for testing"""
    return {
        'efficientnet-b2': {
            'input_size': 260,
            'batch_size': 32,
        },
        'resnet50': {
            'input_size': 224,
            'batch_size': 32,
        }
    }


def create_mock_split_data(num_items=20, split_name="train"):
    """Create mock dataset split data"""
    decades = ['1960s', '1970s', '1980s', '1990s', '2000s']
    mock_data = []

    for i in range(num_items):
        mock_data.append({
            "id": f"{split_name}_id_{i}",
            "product_id": f"product_{i % 5}",  # Some products have multiple images
            "name": f"Test Product {i}",
            "decade": decades[i % 5],  # Ensure all decades represented
            "url": f"https://example.com/{split_name}_image_{i}.jpg",
            "classification": f"furniture",
            "makers": f"Test Maker {i % 3}",
            "country": f"Country_{i % 3}"
        })

    return mock_data


def create_test_environment():
    """Create a complete test environment with mock data"""
    temp_dir = tempfile.mkdtemp()
    temp_path = Path(temp_dir)

    # Create directory structure
    data_dir = temp_path / "data"
    splits_dir = data_dir / "splits"
    cache_dir = data_dir / "cache" / "images"

    splits_dir.mkdir(parents=True)
    cache_dir.mkdir(parents=True)

    # Create split files
    split_files = {}
    for split_name in ["train", "val", "test"]:
        split_data = create_mock_split_data(20 if split_name == "train" else 10, split_name)
        split_file = splits_dir / f"{split_name}.json"

        with open(split_file, 'w') as f:
            json.dump(split_data, f, indent=2)

        split_files[split_name] = split_file

    return temp_path, data_dir, split_files, cache_dir


def test_imports():
    """Test that all required modules can be imported"""
    print("Testing imports...")

    try:
        # Get project root and add to path
        current_dir = Path.cwd()
        src_path = current_dir / "src"
        if src_path.exists():
            sys.path.insert(0, str(current_dir))

        from src.data.data_utils import (
            create_data_loaders, create_weighted_sampler, analyze_dataset_splits,
            create_test_loader, prepare_data_for_training, get_dataset_statistics
        )
        from src.data.url_dataset import BaseDataset, URLDataset, CachedDataset
        from src.data.transforms import get_transforms_for_model

        print("✓ All data_utils imports successful")

        return True

    except ImportError as e:
        print(f"❌ Import failed: {e}")
        return False


def test_create_data_loaders():
    """Test create_data_loaders function"""
    print("Testing create_data_loaders...")

    # Import after ensuring path is set
    from src.data.data_utils import create_data_loaders
    from torchvision import transforms

    temp_path, data_dir, split_files, cache_dir = create_test_environment()

    try:
        # Mock the model configs to avoid import issues
        with patch('src.data.transforms.TRAINING_CONFIGS', create_mock_model_configs()):
            # Create test config
            config = {
                'model_name': 'efficientnet-b2',
                'batch_size': 4,
                'num_workers': 0,  # Use 0 for testing
                'data_dir': str(data_dir),
                'use_cached': False,
                'use_class_weights': False,
                'use_weighted_sampling': False,
                'augmentation_level': 'medium',
                'max_download_retries': 1,
                'download_timeout': 5
            }

            # Mock image downloads to avoid network calls
            mock_image = Image.fromarray(np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8))

            def mock_load_image(self, item):
                return mock_image

            with patch('src.data.url_dataset.URLDataset._load_image', mock_load_image):
                # Test basic data loader creation
                train_loader, val_loader, class_weights, class_names = create_data_loaders(
                    config,
                    data_dir=data_dir,
                    use_subset=True,
                    subset_fraction=0.5  # Use 50% for testing
                )

                print(f"✓ Created DataLoaders successfully")
                print(f"  Train loader: {len(train_loader)} batches")
                print(f"  Val loader: {len(val_loader)} batches")
                print(f"  Class names: {class_names}")

                # Validate results
                assert len(train_loader) > 0, "Train loader should have batches"
                assert len(val_loader) > 0, "Val loader should have batches"
                assert class_names == ['1960s', '1970s', '1980s', '1990s', '2000s'], "Wrong class names"
                assert class_weights is None, "Class weights should be None when disabled"

                # Test loading a batch
                for batch_idx, batch_data in enumerate(train_loader):
                    if len(batch_data) == 3:
                        images, labels, metadata = batch_data
                    else:
                        images, labels = batch_data[0], batch_data[1]
                        metadata = None

                    print(f"✓ Loaded batch: images={images.shape}, labels={labels.shape}")

                    # Validate batch
                    assert images.shape[0] <= config['batch_size'], "Batch size exceeded"
                    assert images.shape[1] == 3, "Wrong number of channels"
                    assert labels.shape[0] == images.shape[0], "Label count mismatch"

                    break  # Just test first batch

    finally:
        # Cleanup
        import shutil
        shutil.rmtree(temp_path)


def test_weighted_sampler():
    """Test create_weighted_sampler function"""
    print("Testing create_weighted_sampler...")

    from src.data.data_utils import create_weighted_sampler
    from src.data.url_dataset import BaseDataset

    temp_path, data_dir, split_files, cache_dir = create_test_environment()

    try:
        # Create dataset with imbalanced data
        imbalanced_data = []
        decades = ['1960s', '1970s', '1980s', '1990s', '2000s']
        # Create imbalanced distribution
        counts = [20, 5, 10, 3, 15]  # Imbalanced

        item_id = 0
        for decade_idx, count in enumerate(counts):
            for i in range(count):
                imbalanced_data.append({
                    "id": f"id_{item_id}",
                    "product_id": f"product_{item_id}",
                    "name": f"Product {item_id}",
                    "decade": decades[decade_idx],
                    "url": f"https://example.com/image_{item_id}.jpg",
                    "classification": "test",
                    "makers": "test",
                    "country": "test"
                })
                item_id += 1

        # Save imbalanced split
        imbalanced_split = data_dir / "splits" / "imbalanced.json"
        with open(imbalanced_split, 'w') as f:
            json.dump(imbalanced_data, f)

        # Create dataset
        dataset = BaseDataset(str(imbalanced_split))

        print(f"✓ Created imbalanced dataset with {len(dataset)} samples")

        # Test label distribution before sampling
        labels = dataset.get_labels()
        from collections import Counter
        original_counts = Counter(labels)
        print(f"  Original distribution: {dict(original_counts)}")

        # Create weighted sampler
        sampler = create_weighted_sampler(dataset)

        print(f"✓ Created weighted sampler with {len(sampler)} samples")

        # Test sampling
        sampled_indices = list(sampler)[:100]  # Sample 100 items
        sampled_labels = [labels[idx] for idx in sampled_indices]
        sampled_counts = Counter(sampled_labels)

        print(f"  Sampled distribution: {dict(sampled_counts)}")

        # Check if sampling is more balanced
        original_ratio = max(original_counts.values()) / min(original_counts.values())
        sampled_ratio = max(sampled_counts.values()) / min(sampled_counts.values()) if min(
            sampled_counts.values()) > 0 else float('inf')

        print(f"  Balance improvement: {original_ratio:.2f} -> {sampled_ratio:.2f}")

        # Sampled distribution should be more balanced
        assert sampled_ratio < original_ratio or sampled_ratio < 3.0, "Sampling should improve balance"

    finally:
        import shutil
        shutil.rmtree(temp_path)


def test_analyze_dataset_splits():
    """Test analyze_dataset_splits function"""
    print("Testing analyze_dataset_splits...")

    from src.data.data_utils import analyze_dataset_splits

    temp_path, data_dir, split_files, cache_dir = create_test_environment()

    try:
        # Analyze the splits
        analysis = analyze_dataset_splits(data_dir)

        print(f"✓ Analysis completed for {len(analysis)} splits")

        # Check analysis structure
        expected_splits = ['train', 'val', 'test']
        for split_name in expected_splits:
            assert split_name in analysis, f"Missing analysis for {split_name}"

            split_stats = analysis[split_name]
            print(f"  {split_name}: {split_stats['total_images']} images, {split_stats['unique_products']} products")

            # Validate structure
            required_keys = ['total_images', 'unique_products', 'decades', 'classifications', 'countries']
            for key in required_keys:
                assert key in split_stats, f"Missing key {key} in {split_name} analysis"

            # Check decades distribution
            decades_dist = split_stats['decades']
            assert len(decades_dist) == 5, f"Should have 5 decades, got {len(decades_dist)}"

        # Check for data leakage detection
        # Our test data has overlapping product_ids, so leakage should be detected
        if 'data_leakage' in analysis:
            print(f"  ⚠️  Data leakage detected: {analysis['data_leakage']['train_val_overlap']} overlapping products")
        else:
            print(f"  ✓ No data leakage detected")

    finally:
        import shutil
        shutil.rmtree(temp_path)


def test_prepare_data_for_training():
    """Test prepare_data_for_training function"""
    print("Testing prepare_data_for_training...")

    from src.data.data_utils import prepare_data_for_training

    temp_path, data_dir, split_files, cache_dir = create_test_environment()

    try:
        # Test config
        config = {
            'data_dir': str(data_dir),
            'use_cached': False
        }

        # FIXED: Test with data leakage detection disabled first
        is_ready = prepare_data_for_training(
            config,
            download_if_missing=False,
            verify_splits=False  # Disable verification to test basic functionality
        )

        print(f"✓ Data preparation (no verification): {is_ready}")
        assert is_ready == True, "Data preparation should succeed when verification disabled"

        # FIXED: Test with verification enabled (should detect leakage)
        is_ready_with_verification = prepare_data_for_training(
            config,
            download_if_missing=False,
            verify_splits=True  # Enable verification
        )

        print(f"✓ Data preparation (with verification): {is_ready_with_verification}")
        # This should be False because our mock data has intentional leakage
        assert is_ready_with_verification == False, "Data preparation should fail when data leakage detected"

        # FIXED: Test with missing split files
        missing_split = data_dir / "splits" / "train.json"
        missing_split.unlink()  # Remove train split

        is_ready_missing = prepare_data_for_training(
            config,
            download_if_missing=False,
            verify_splits=False
        )

        print(f"✓ Data preparation with missing files: {is_ready_missing}")
        assert is_ready_missing == False, "Data preparation should fail with missing splits"

        print("✓ All data preparation scenarios tested correctly")

    finally:
        import shutil
        shutil.rmtree(temp_path)


def test_get_dataset_statistics():
    """Test get_dataset_statistics function"""
    print("Testing get_dataset_statistics...")

    from src.data.data_utils import get_dataset_statistics
    from src.data.url_dataset import BaseDataset

    temp_path, data_dir, split_files, cache_dir = create_test_environment()

    try:
        # Create dataset
        dataset = BaseDataset(str(split_files['train']))

        # Get statistics
        stats = get_dataset_statistics(dataset)

        print(f"✓ Generated statistics for dataset")

        # Validate statistics structure
        required_keys = [
            'total_samples', 'num_classes', 'class_names',
            'class_distribution', 'class_balance', 'product_stats',
            'top_classifications', 'top_countries'
        ]

        for key in required_keys:
            assert key in stats, f"Missing statistics key: {key}"

        print(f"  Total samples: {stats['total_samples']}")
        print(f"  Number of classes: {stats['num_classes']}")
        print(f"  Class distribution: {stats['class_distribution']}")
        print(f"  Imbalance ratio: {stats['class_balance']['imbalance_ratio']:.2f}")
        print(f"  Unique products: {stats['product_stats']['unique_products']}")
        print(f"  Avg images per product: {stats['product_stats']['avg_images_per_product']:.2f}")

        # Validate values
        assert stats['total_samples'] == len(dataset), "Wrong total samples count"
        assert stats['num_classes'] == 5, "Wrong number of classes"
        assert len(stats['class_names']) == 5, "Wrong number of class names"

    finally:
        import shutil
        shutil.rmtree(temp_path)


def test_create_test_loader():
    """Test create_test_loader function"""
    print("Testing create_test_loader...")

    from src.data.data_utils import create_test_loader

    temp_path, data_dir, split_files, cache_dir = create_test_environment()

    try:
        # Mock the model configs
        with patch('src.data.transforms.TRAINING_CONFIGS', create_mock_model_configs()):
            config = {
                'model_name': 'resnet50',
                'batch_size': 4,
                'num_workers': 0,
                'data_dir': str(data_dir),
                'use_cached': False,
                'max_download_retries': 1,
                'download_timeout': 5
            }

            # Mock image loading
            mock_image = Image.fromarray(np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8))

            def mock_load_image(self, item):
                return mock_image

            with patch('src.data.url_dataset.URLDataset._load_image', mock_load_image):
                test_loader, class_names = create_test_loader(
                    config,
                    data_dir=data_dir,
                    batch_size=6
                )

                print(f"✓ Created test loader with {len(test_loader)} batches")
                print(f"  Class names: {class_names}")

                # Validate
                assert len(test_loader) > 0, "Test loader should have batches"
                assert class_names == ['1960s', '1970s', '1980s', '1990s', '2000s'], "Wrong class names"

                # Test loading a batch
                for batch_data in test_loader:
                    if len(batch_data) == 3:
                        images, labels, metadata = batch_data
                    else:
                        images, labels = batch_data[0], batch_data[1]

                    print(f"✓ Test batch: images={images.shape}, labels={labels.shape}")
                    assert images.shape[0] <= 6, "Batch size should respect config"
                    break

    finally:
        import shutil
        shutil.rmtree(temp_path)


def test_integration_with_real_data():
    """Test integration with real project data if available"""
    print("Testing integration with real data...")

    # Check if real data exists
    project_root = Path.cwd()
    real_data_dir = project_root / "data"
    real_splits_dir = real_data_dir / "splits"
    real_train_split = real_splits_dir / "train.json"

    if not real_train_split.exists():
        print("⚠️  Real data not found, skipping integration test")
        raise FileNotFoundError("Real data not available")

    from src.data.data_utils import create_data_loaders, analyze_dataset_splits

    try:
        # Mock the model configs
        with patch('src.data.transforms.TRAINING_CONFIGS', create_mock_model_configs()):
            print(f"✓ Found real data at {real_data_dir}")

            # Test analysis with real data
            analysis = analyze_dataset_splits(real_data_dir)

            for split_name, stats in analysis.items():
                if isinstance(stats, dict) and 'total_images' in stats:
                    print(f"  {split_name}: {stats['total_images']} images")

            # Test data loader creation with real data (small subset)
            config = {
                'model_name': 'efficientnet-b2',
                'batch_size': 4,
                'num_workers': 0,
                'data_dir': str(real_data_dir),
                'use_cached': False,
                'use_class_weights': False,
                'use_weighted_sampling': False,
                'augmentation_level': 'light',
                'max_download_retries': 1,
                'download_timeout': 5
            }

            # Use very small subset for quick testing
            train_loader, val_loader, class_weights, class_names = create_data_loaders(
                config,
                data_dir=real_data_dir,
                use_subset=True,
                subset_fraction=0.001  # 0.1% for very quick test
            )

            print(f"✓ Real data integration successful")
            print(f"  Train batches: {len(train_loader)}")
            print(f"  Val batches: {len(val_loader)}")

    except Exception as e:
        print(f"⚠️  Real data integration test failed: {e}")
        raise


def main():
    """Main test function"""
    print("DATA_UTILS.PY COMPREHENSIVE TEST SUITE")
    print("=" * 60)
    print("Testing data orchestration and integration layer")

    runner = TestRunner()

    # Check imports first
    if not test_imports():
        print("❌ Cannot proceed without successful imports")
        return 1

    # Define all tests
    tests = [
        ("Create Data Loaders", test_create_data_loaders),
        ("Weighted Sampler", test_weighted_sampler),
        ("Analyze Dataset Splits", test_analyze_dataset_splits),
        ("Prepare Data for Training", test_prepare_data_for_training),
        ("Get Dataset Statistics", test_get_dataset_statistics),
        ("Create Test Loader", test_create_test_loader),
        ("Integration with Real Data", test_integration_with_real_data),
    ]

    # Run all tests
    for test_name, test_func in tests:
        runner.run_test(test_name, test_func)

    # Print summary
    runner.print_summary()

    # Return exit code
    success = runner.failed == 0
    if success:
        print("\n🎉 ALL DATA_UTILS TESTS PASSED!")
        print("✅ The complete data pipeline is working correctly!")
        print("\nYour data loading system is ready for:")
        print("  🚀 Model training")
        print("  📊 Data analysis")
        print("  🔄 Production deployment")
    else:
        print("\n❌ SOME TESTS FAILED!")
        print("Check the issues above before proceeding to training")

    return 0 if success else 1


if __name__ == "__main__":
    exit_code = main()
    print(f"\nTest completed with exit code: {exit_code}")


	# !/usr/bin/env python3
	"""
	Comprehensive test script for data_utils.py
	Tests the orchestration layer that combines transforms, url_dataset, and creates DataLoaders
	Run from project root: python test_data_utils_comprehensive.py
	"""

	import sys
	import os
	from pathlib import Path
	import torch
	import numpy as np
	from PIL import Image
	import json
	import tempfile
	import traceback
	import logging
	from typing import Dict, List
	from unittest.mock import patch, MagicMock

	# Set up logging
	logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')


	class TestRunner:
	"""Test runner with results tracking"""

	def __init__(self):
	self.passed = 0
	self.failed = 0
	self.skipped = 0
	self.errors = []

	def run_test(self, test_name: str, test_func):
	"""Run a single test with error handling"""
	print(f"\n{'=' * 60}")
	print(f"Running: {test_name}")
	print(f"{'=' * 60}")

	try:
	test_func()
	self.passed += 1
	print(f"✅ PASSED: {test_name}")
	except FileNotFoundError as e:
	self.skipped += 1
	print(f"⚠️ SKIPPED: {test_name} - {str(e)}")
	except Exception as e:
	self.failed += 1
	self.errors.append({
	'test': test_name,
	'error': str(e),
	'traceback': traceback.format_exc()
	})
	print(f"❌ FAILED: {test_name}")
	print(f"Error: {str(e)}")
	print(f"Traceback:\n{traceback.format_exc()}")

	def print_summary(self):
	"""Print test summary"""
	total = self.passed + self.failed + self.skipped
	print(f"\n{'=' * 60}")
	print(f"TEST SUMMARY")
	print(f"{'=' * 60}")
	print(f"Total tests: {total}")
	print(f"✅ Passed: {self.passed}")
	print(f"❌ Failed: {self.failed}")
	print(f"⚠️ Skipped: {self.skipped}")

	if self.errors:
	print(f"\nFailed tests:")
	for error in self.errors:
	print(f" - {error['test']}: {error['error']}")


	def create_mock_model_configs():
	"""Create mock model configs for testing"""
	return {
	'efficientnet-b2': {
	'input_size': 260,
	'batch_size': 32,
	},
	'resnet50': {
	'input_size': 224,
	'batch_size': 32,
	}
	}


	def create_mock_split_data(num_items=20, split_name="train"):
	"""Create mock dataset split data"""
	decades = ['1960s', '1970s', '1980s', '1990s', '2000s']
	mock_data = []

	for i in range(num_items):
	mock_data.append({
	"id": f"{split_name}_id_{i}",
	"product_id": f"product_{i % 5}", # Some products have multiple images
	"name": f"Test Product {i}",
	"decade": decades[i % 5], # Ensure all decades represented
	"url": f"https://example.com/{split_name}_image_{i}.jpg",
	"classification": f"furniture",
	"makers": f"Test Maker {i % 3}",
	"country": f"Country_{i % 3}"
	})

	return mock_data


	def create_test_environment():
	"""Create a complete test environment with mock data"""
	temp_dir = tempfile.mkdtemp()
	temp_path = Path(temp_dir)

	# Create directory structure
	data_dir = temp_path / "data"
	splits_dir = data_dir / "splits"
	cache_dir = data_dir / "cache" / "images"

	splits_dir.mkdir(parents=True)
	cache_dir.mkdir(parents=True)

	# Create split files
	split_files = {}
	for split_name in ["train", "val", "test"]:
	split_data = create_mock_split_data(20 if split_name == "train" else 10, split_name)
	split_file = splits_dir / f"{split_name}.json"

	with open(split_file, 'w') as f:
	json.dump(split_data, f, indent=2)

	split_files[split_name] = split_file

	return temp_path, data_dir, split_files, cache_dir


	def test_imports():
	"""Test that all required modules can be imported"""
	print("Testing imports...")

	try:
	# Get project root and add to path
	current_dir = Path.cwd()
	src_path = current_dir / "src"
	if src_path.exists():
	sys.path.insert(0, str(current_dir))

	from src.data.data_utils import (
	create_data_loaders, create_weighted_sampler, analyze_dataset_splits,
	create_test_loader, prepare_data_for_training, get_dataset_statistics
	)
	from src.data.url_dataset import BaseDataset, URLDataset, CachedDataset
	from src.data.transforms import get_transforms_for_model

	print("✓ All data_utils imports successful")

	return True

	except ImportError as e:
	print(f"❌ Import failed: {e}")
	return False


	def test_create_data_loaders():
	"""Test create_data_loaders function"""
	print("Testing create_data_loaders...")

	# Import after ensuring path is set
	from src.data.data_utils import create_data_loaders
	from torchvision import transforms

	temp_path, data_dir, split_files, cache_dir = create_test_environment()

	try:
	# Mock the model configs to avoid import issues
	with patch('src.data.transforms.TRAINING_CONFIGS', create_mock_model_configs()):
	# Create test config
	config = {
	'model_name': 'efficientnet-b2',
	'batch_size': 4,
	'num_workers': 0, # Use 0 for testing
	'data_dir': str(data_dir),
	'use_cached': False,
	'use_class_weights': False,
	'use_weighted_sampling': False,
	'augmentation_level': 'medium',
	'max_download_retries': 1,
	'download_timeout': 5
	}

	# Mock image downloads to avoid network calls
	mock_image = Image.fromarray(np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8))

	def mock_load_image(self, item):
	return mock_image

	with patch('src.data.url_dataset.URLDataset._load_image', mock_load_image):
	# Test basic data loader creation
	train_loader, val_loader, class_weights, class_names = create_data_loaders(
	config,
	data_dir=data_dir,
	use_subset=True,
	subset_fraction=0.5 # Use 50% for testing
	)

	print(f"✓ Created DataLoaders successfully")
	print(f" Train loader: {len(train_loader)} batches")
	print(f" Val loader: {len(val_loader)} batches")
	print(f" Class names: {class_names}")

	# Validate results
	assert len(train_loader) > 0, "Train loader should have batches"
	assert len(val_loader) > 0, "Val loader should have batches"
	assert class_names == ['1960s', '1970s', '1980s', '1990s', '2000s'], "Wrong class names"
	assert class_weights is None, "Class weights should be None when disabled"

	# Test loading a batch
	for batch_idx, batch_data in enumerate(train_loader):
	if len(batch_data) == 3:
	images, labels, metadata = batch_data
	else:
	images, labels = batch_data[0], batch_data[1]
	metadata = None

	print(f"✓ Loaded batch: images={images.shape}, labels={labels.shape}")

	# Validate batch
	assert images.shape[0] <= config['batch_size'], "Batch size exceeded"
	assert images.shape[1] == 3, "Wrong number of channels"
	assert labels.shape[0] == images.shape[0], "Label count mismatch"

	break # Just test first batch

	finally:
	# Cleanup
	import shutil
	shutil.rmtree(temp_path)


	def test_weighted_sampler():
	"""Test create_weighted_sampler function"""
	print("Testing create_weighted_sampler...")

	from src.data.data_utils import create_weighted_sampler
	from src.data.url_dataset import BaseDataset

	temp_path, data_dir, split_files, cache_dir = create_test_environment()

	try:
	# Create dataset with imbalanced data
	imbalanced_data = []
	decades = ['1960s', '1970s', '1980s', '1990s', '2000s']
	# Create imbalanced distribution
	counts = [20, 5, 10, 3, 15] # Imbalanced

	item_id = 0
	for decade_idx, count in enumerate(counts):
	for i in range(count):
	imbalanced_data.append({
	"id": f"id_{item_id}",
	"product_id": f"product_{item_id}",
	"name": f"Product {item_id}",
	"decade": decades[decade_idx],
	"url": f"https://example.com/image_{item_id}.jpg",
	"classification": "test",
	"makers": "test",
	"country": "test"
	})
	item_id += 1

	# Save imbalanced split
	imbalanced_split = data_dir / "splits" / "imbalanced.json"
	with open(imbalanced_split, 'w') as f:
	json.dump(imbalanced_data, f)

	# Create dataset
	dataset = BaseDataset(str(imbalanced_split))

	print(f"✓ Created imbalanced dataset with {len(dataset)} samples")

	# Test label distribution before sampling
	labels = dataset.get_labels()
	from collections import Counter
	original_counts = Counter(labels)
	print(f" Original distribution: {dict(original_counts)}")

	# Create weighted sampler
	sampler = create_weighted_sampler(dataset)

	print(f"✓ Created weighted sampler with {len(sampler)} samples")

	# Test sampling
	sampled_indices = list(sampler)[:100] # Sample 100 items
	sampled_labels = [labels[idx] for idx in sampled_indices]
	sampled_counts = Counter(sampled_labels)

	print(f" Sampled distribution: {dict(sampled_counts)}")

	# Check if sampling is more balanced
	original_ratio = max(original_counts.values()) / min(original_counts.values())
	sampled_ratio = max(sampled_counts.values()) / min(sampled_counts.values()) if min(
	sampled_counts.values()) > 0 else float('inf')

	print(f" Balance improvement: {original_ratio:.2f} -> {sampled_ratio:.2f}")

	# Sampled distribution should be more balanced
	assert sampled_ratio < original_ratio or sampled_ratio < 3.0, "Sampling should improve balance"

	finally:
	import shutil
	shutil.rmtree(temp_path)


	def test_analyze_dataset_splits():
	"""Test analyze_dataset_splits function"""
	print("Testing analyze_dataset_splits...")

	from src.data.data_utils import analyze_dataset_splits

	temp_path, data_dir, split_files, cache_dir = create_test_environment()

	try:
	# Analyze the splits
	analysis = analyze_dataset_splits(data_dir)

	print(f"✓ Analysis completed for {len(analysis)} splits")

	# Check analysis structure
	expected_splits = ['train', 'val', 'test']
	for split_name in expected_splits:
	assert split_name in analysis, f"Missing analysis for {split_name}"

	split_stats = analysis[split_name]
	print(f" {split_name}: {split_stats['total_images']} images, {split_stats['unique_products']} products")

	# Validate structure
	required_keys = ['total_images', 'unique_products', 'decades', 'classifications', 'countries']
	for key in required_keys:
	assert key in split_stats, f"Missing key {key} in {split_name} analysis"

	# Check decades distribution
	decades_dist = split_stats['decades']
	assert len(decades_dist) == 5, f"Should have 5 decades, got {len(decades_dist)}"

	# Check for data leakage detection
	# Our test data has overlapping product_ids, so leakage should be detected
	if 'data_leakage' in analysis:
	print(f" ⚠️ Data leakage detected: {analysis['data_leakage']['train_val_overlap']} overlapping products")
	else:
	print(f" ✓ No data leakage detected")

	finally:
	import shutil
	shutil.rmtree(temp_path)


	def test_prepare_data_for_training():
	"""Test prepare_data_for_training function"""
	print("Testing prepare_data_for_training...")

	from src.data.data_utils import prepare_data_for_training

	temp_path, data_dir, split_files, cache_dir = create_test_environment()

	try:
	# Test config
	config = {
	'data_dir': str(data_dir),
	'use_cached': False
	}

	# FIXED: Test with data leakage detection disabled first
	is_ready = prepare_data_for_training(
	config,
	download_if_missing=False,
	verify_splits=False # Disable verification to test basic functionality
	)

	print(f"✓ Data preparation (no verification): {is_ready}")
	assert is_ready == True, "Data preparation should succeed when verification disabled"

	# FIXED: Test with verification enabled (should detect leakage)
	is_ready_with_verification = prepare_data_for_training(
	config,
	download_if_missing=False,
	verify_splits=True # Enable verification
	)

	print(f"✓ Data preparation (with verification): {is_ready_with_verification}")
	# This should be False because our mock data has intentional leakage
	assert is_ready_with_verification == False, "Data preparation should fail when data leakage detected"

	# FIXED: Test with missing split files
	missing_split = data_dir / "splits" / "train.json"
	missing_split.unlink() # Remove train split

	is_ready_missing = prepare_data_for_training(
	config,
	download_if_missing=False,
	verify_splits=False
	)

	print(f"✓ Data preparation with missing files: {is_ready_missing}")
	assert is_ready_missing == False, "Data preparation should fail with missing splits"

	print("✓ All data preparation scenarios tested correctly")

	finally:
	import shutil
	shutil.rmtree(temp_path)


	def test_get_dataset_statistics():
	"""Test get_dataset_statistics function"""
	print("Testing get_dataset_statistics...")

	from src.data.data_utils import get_dataset_statistics
	from src.data.url_dataset import BaseDataset

	temp_path, data_dir, split_files, cache_dir = create_test_environment()

	try:
	# Create dataset
	dataset = BaseDataset(str(split_files['train']))

	# Get statistics
	stats = get_dataset_statistics(dataset)

	print(f"✓ Generated statistics for dataset")

	# Validate statistics structure
	required_keys = [
	'total_samples', 'num_classes', 'class_names',
	'class_distribution', 'class_balance', 'product_stats',
	'top_classifications', 'top_countries'
	]

	for key in required_keys:
	assert key in stats, f"Missing statistics key: {key}"

	print(f" Total samples: {stats['total_samples']}")
	print(f" Number of classes: {stats['num_classes']}")
	print(f" Class distribution: {stats['class_distribution']}")
	print(f" Imbalance ratio: {stats['class_balance']['imbalance_ratio']:.2f}")
	print(f" Unique products: {stats['product_stats']['unique_products']}")
	print(f" Avg images per product: {stats['product_stats']['avg_images_per_product']:.2f}")

	# Validate values
	assert stats['total_samples'] == len(dataset), "Wrong total samples count"
	assert stats['num_classes'] == 5, "Wrong number of classes"
	assert len(stats['class_names']) == 5, "Wrong number of class names"

	finally:
	import shutil
	shutil.rmtree(temp_path)


	def test_create_test_loader():
	"""Test create_test_loader function"""
	print("Testing create_test_loader...")

	from src.data.data_utils import create_test_loader

	temp_path, data_dir, split_files, cache_dir = create_test_environment()

	try:
	# Mock the model configs
	with patch('src.data.transforms.TRAINING_CONFIGS', create_mock_model_configs()):
	config = {
	'model_name': 'resnet50',
	'batch_size': 4,
	'num_workers': 0,
	'data_dir': str(data_dir),
	'use_cached': False,
	'max_download_retries': 1,
	'download_timeout': 5
	}

	# Mock image loading
	mock_image = Image.fromarray(np.random.randint(0, 255, (200, 200, 3), dtype=np.uint8))

	def mock_load_image(self, item):
	return mock_image

	with patch('src.data.url_dataset.URLDataset._load_image', mock_load_image):
	test_loader, class_names = create_test_loader(
	config,
	data_dir=data_dir,
	batch_size=6
	)

	print(f"✓ Created test loader with {len(test_loader)} batches")
	print(f" Class names: {class_names}")

	# Validate
	assert len(test_loader) > 0, "Test loader should have batches"
	assert class_names == ['1960s', '1970s', '1980s', '1990s', '2000s'], "Wrong class names"

	# Test loading a batch
	for batch_data in test_loader:
	if len(batch_data) == 3:
	images, labels, metadata = batch_data
	else:
	images, labels = batch_data[0], batch_data[1]

	print(f"✓ Test batch: images={images.shape}, labels={labels.shape}")
	assert images.shape[0] <= 6, "Batch size should respect config"
	break

	finally:
	import shutil
	shutil.rmtree(temp_path)


	def test_integration_with_real_data():
	"""Test integration with real project data if available"""
	print("Testing integration with real data...")

	# Check if real data exists
	project_root = Path.cwd()
	real_data_dir = project_root / "data"
	real_splits_dir = real_data_dir / "splits"
	real_train_split = real_splits_dir / "train.json"

	if not real_train_split.exists():
	print("⚠️ Real data not found, skipping integration test")
	raise FileNotFoundError("Real data not available")

	from src.data.data_utils import create_data_loaders, analyze_dataset_splits

	try:
	# Mock the model configs
	with patch('src.data.transforms.TRAINING_CONFIGS', create_mock_model_configs()):
	print(f"✓ Found real data at {real_data_dir}")

	# Test analysis with real data
	analysis = analyze_dataset_splits(real_data_dir)

	for split_name, stats in analysis.items():
	if isinstance(stats, dict) and 'total_images' in stats:
	print(f" {split_name}: {stats['total_images']} images")

	# Test data loader creation with real data (small subset)
	config = {
	'model_name': 'efficientnet-b2',
	'batch_size': 4,
	'num_workers': 0,
	'data_dir': str(real_data_dir),
	'use_cached': False,
	'use_class_weights': False,
	'use_weighted_sampling': False,
	'augmentation_level': 'light',
	'max_download_retries': 1,
	'download_timeout': 5
	}

	# Use very small subset for quick testing
	train_loader, val_loader, class_weights, class_names = create_data_loaders(
	config,
	data_dir=real_data_dir,
	use_subset=True,
	subset_fraction=0.001 # 0.1% for very quick test
	)

	print(f"✓ Real data integration successful")
	print(f" Train batches: {len(train_loader)}")
	print(f" Val batches: {len(val_loader)}")

	except Exception as e:
	print(f"⚠️ Real data integration test failed: {e}")
	raise


	def main():
	"""Main test function"""
	print("DATA_UTILS.PY COMPREHENSIVE TEST SUITE")
	print("=" * 60)
	print("Testing data orchestration and integration layer")

	runner = TestRunner()

	# Check imports first
	if not test_imports():
	print("❌ Cannot proceed without successful imports")
	return 1

	# Define all tests
	tests = [
	("Create Data Loaders", test_create_data_loaders),
	("Weighted Sampler", test_weighted_sampler),
	("Analyze Dataset Splits", test_analyze_dataset_splits),
	("Prepare Data for Training", test_prepare_data_for_training),
	("Get Dataset Statistics", test_get_dataset_statistics),
	("Create Test Loader", test_create_test_loader),
	("Integration with Real Data", test_integration_with_real_data),
	]

	# Run all tests
	for test_name, test_func in tests:
	runner.run_test(test_name, test_func)

	# Print summary
	runner.print_summary()

	# Return exit code
	success = runner.failed == 0
	if success:
	print("\n🎉 ALL DATA_UTILS TESTS PASSED!")
	print("✅ The complete data pipeline is working correctly!")
	print("\nYour data loading system is ready for:")
	print(" 🚀 Model training")
	print(" 📊 Data analysis")
	print(" 🔄 Production deployment")
	else:
	print("\n❌ SOME TESTS FAILED!")
	print("Check the issues above before proceeding to training")

	return 0 if success else 1


	if __name__ == "__main__":
	exit_code = main()
	print(f"\nTest completed with exit code: {exit_code}")