test_data_pipeline.py

"""
Test module for data pipeline functionality.
"""

import json
import tempfile
import unittest.mock
from pathlib import Path
from unittest.mock import MagicMock, patch

import pandas as pd
import pytest
import yaml


class TestDataPipelineIntegration:
    """Test data pipeline integration functionality."""

    def test_end_to_end_parquet_dataset(self, tmp_path):
        """Test end-to-end parquet dataset creation and loading."""
        # Create test data
        test_data = pd.DataFrame({
            'id': ['1', '2', '3'],
            'question': ['What is AI?', 'What is ML?', 'What is DL?'],
            'choices': [['A', 'B', 'C'], ['X', 'Y', 'Z'], ['1', '2', '3']],
            'answer': ['A', 'X', '1']
        })
        
        # Test parquet operations
        parquet_file = tmp_path / "test_data.parquet"
        test_data.to_parquet(parquet_file)
        
        # Verify file was created and can be read
        assert parquet_file.exists()
        loaded_data = pd.read_parquet(parquet_file)
        assert len(loaded_data) == 3
        assert list(loaded_data.columns) == ['id', 'question', 'choices', 'answer']

    def test_makefile_integration(self):
        """Test Makefile integration."""
        # This test verifies that the Makefile targets work
        # In a real implementation, this would test actual Makefile commands
        assert True  # Placeholder for now

    def test_schema_validation(self, tmp_path):
        """Test schema validation for parquet files."""
        # Create test data with specific schema
        test_data = pd.DataFrame({
            'id': ['test_1', 'test_2'],
            'content': ['sample content 1', 'sample content 2'],
            'metadata': [{'type': 'test'}, {'type': 'test'}]
        })
        
        # Test parquet operations with schema validation
        parquet_file = tmp_path / "schema_test.parquet"
        test_data.to_parquet(parquet_file)
        
        # Verify schema
        loaded_data = pd.read_parquet(parquet_file)
        assert 'id' in loaded_data.columns
        assert 'content' in loaded_data.columns
        assert 'metadata' in loaded_data.columns


class TestDataPipelineErrorHandling:
    """Test error handling in data pipeline."""

    def test_network_error_handling(self):
        """Test handling of network errors during data fetching."""
        # Mock network error
        with patch('requests.get') as mock_get:
            mock_get.side_effect = ConnectionError("Network error")
            
            # Test that network errors are handled gracefully
            # In a real implementation, this would test actual error handling
            assert True  # Placeholder for now

    def test_checksum_mismatch_recovery(self, tmp_path):
        """Test recovery from checksum mismatches."""
        # Create mock download function with correct signature
        def mock_download(url, destination, description=None):
            """Mock download function that takes 3 arguments."""
            # Create a dummy file
            Path(destination).parent.mkdir(parents=True, exist_ok=True)
            Path(destination).write_text("dummy content")
            return True
        
        # Mock the fetch_data module
        with patch('scripts.fetch_data.download_with_progress', side_effect=mock_download):
            with patch('scripts.fetch_data.fetch_dataset') as mock_fetch:
                mock_fetch.return_value = True
                
                # Test checksum mismatch recovery
                result = mock_fetch("test")
                assert result is True

    def test_adapter_failure_handling(self):
        """Test handling of adapter failures."""
        # Mock adapter failure using a proper module path
        with patch('scripts.fetch_data.fetch_dataset') as mock_adapter:
            mock_adapter.side_effect = Exception("Adapter failure")
            
            # Test that adapter failures are handled gracefully
            try:
                mock_adapter("test_dataset")
                assert False, "Should have raised an exception"
            except Exception as e:
                assert str(e) == "Adapter failure"


class TestCLIIntegration:
    """Test CLI integration functionality."""

    def test_cli_list_datasets(self, tmp_path):
        """Test CLI list datasets functionality."""
        # Create test registry structure
        registry_dir = tmp_path / "data"
        registry_dir.mkdir(parents=True, exist_ok=True)
        
        # Create registry file (not directory) - ensure parent directory exists
        registry_file = registry_dir / "registry.yaml"
        
        # Ensure the parent directory exists and is not a file
        if registry_file.exists() and registry_file.is_file():
            registry_file.unlink()  # Remove if it's a file
        
        registry_data = {
            'datasets': {
                'test_dataset': {
                    'url': 'https://example.com/test.tar.gz',
                    'checksum': 'test_checksum'
                }
            }
        }
        
        # Write registry file
        registry_file.write_text(yaml.dump(registry_data))
        
        # Verify registry file exists and is readable
        assert registry_file.exists()
        assert registry_file.is_file()
        
        # Test reading registry
        loaded_data = yaml.safe_load(registry_file.read_text())
        assert 'datasets' in loaded_data
        assert 'test_dataset' in loaded_data['datasets']

    def test_cli_compute_checksum(self):
        """Test CLI compute checksum functionality."""
        # Test checksum computation
        test_content = "test content for checksum"
        # In a real implementation, this would compute actual checksums
        assert len(test_content) > 0

    def test_cli_force_redownload(self, tmp_path):
        """Test CLI force redownload functionality."""
        # Create mock download function with correct signature
        def mock_download(url, destination, description=None):
            """Mock download function that takes 3 arguments."""
            # Create a dummy file
            Path(destination).parent.mkdir(parents=True, exist_ok=True)
            Path(destination).write_text("dummy content")
            return True
        
        # Mock the fetch_data module
        with patch('scripts.fetch_data.download_with_progress', side_effect=mock_download):
            with patch('scripts.fetch_data.fetch_dataset') as mock_fetch:
                mock_fetch.return_value = True
                
                # Test force redownload
                result = mock_fetch("test", force=True)
                assert result is True