Team Ai
Modelpublic

ParallelLLC/algorithmic_trading

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
27likes32downloads
test_data_ingestion.py338 linesDownload Raw Back to tests
1import pytest2import pandas as pd3import numpy as np4import tempfile5import os6from unittest.mock import patch, MagicMock7from agentic_ai_system.data_ingestion import load_data, validate_data, _load_csv_data, _load_synthetic_data8 9class TestDataIngestion:10    """Test cases for data ingestion module"""11    12    @pytest.fixture13    def config(self):14        """Sample configuration for testing"""15        return {16            'data_source': {17                'type': 'csv',18                'path': 'data/market_data.csv'19            },20            'synthetic_data': {21                'base_price': 150.0,22                'volatility': 0.02,23                'trend': 0.001,24                'noise_level': 0.005,25                'data_path': 'data/synthetic_market_data.csv'26            },27            'trading': {28                'symbol': 'AAPL',29                'timeframe': '1min'30            }31        }32    33    @pytest.fixture34    def sample_csv_data(self):35        """Create sample CSV data for testing"""36        dates = pd.date_range(start='2024-01-01', periods=100, freq='1min')37        38        data = []39        for i, date in enumerate(dates):40            base_price = 150.0 + (i * 0.1)41            42            # Generate OHLC values that follow proper relationships43            open_price = base_price + np.random.normal(0, 1)44            close_price = base_price + np.random.normal(0, 1)45            46            # High should be >= max(open, close)47            high_price = max(open_price, close_price) + abs(np.random.normal(0, 1))48            49            # Low should be <= min(open, close)50            low_price = min(open_price, close_price) - abs(np.random.normal(0, 1))51            52            data.append({53                'timestamp': date,54                'open': open_price,55                'high': high_price,56                'low': low_price,57                'close': close_price,58                'volume': np.random.randint(1000, 100000)59            })60        61        return pd.DataFrame(data)62    63    def test_load_data_csv_type(self, config, sample_csv_data):64        """Test loading data with CSV type"""65        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as tmp_file:66            sample_csv_data.to_csv(tmp_file.name, index=False)67            config['data_source']['path'] = tmp_file.name68            69            try:70                result = load_data(config)71                72                assert isinstance(result, pd.DataFrame)73                assert len(result) == len(sample_csv_data)74                assert list(result.columns) == list(sample_csv_data.columns)75                76            finally:77                os.unlink(tmp_file.name)78    79    def test_load_data_synthetic_type(self, config):80        """Test loading data with synthetic type"""81        config['data_source']['type'] = 'synthetic'82        83        with patch('agentic_ai_system.data_ingestion._load_synthetic_data') as mock_generate:84            mock_df = pd.DataFrame({85                'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),86                'open': [150] * 10,87                'high': [155] * 10,88                'low': [145] * 10,89                'close': [152] * 10,90                'volume': [1000] * 1091            })92            mock_generate.return_value = mock_df93            94            result = load_data(config)95            96            assert isinstance(result, pd.DataFrame)97            mock_generate.assert_called_once_with(config)98 99    def test_load_data_yahoo_type(self, config):100        config['data_source']['type'] = 'yahoo'101        mock_df = pd.DataFrame({102            'timestamp': pd.date_range('2024-01-01', periods=10, freq='D'),103            'open': [150] * 10,104            'high': [155] * 10,105            'low': [145] * 10,106            'close': [152] * 10,107            'volume': [1000] * 10,108        })109        mock_stream = MagicMock()110        mock_stream.get_historical_data.return_value = mock_df111        with patch('agentic_ai_system.yahoo_data_stream.YahooDataStream', return_value=mock_stream):112            result = load_data(config)113        assert isinstance(result, pd.DataFrame)114        assert len(result) == 10115        mock_stream.get_historical_data.assert_called_once()116 117    def test_load_data_invalid_type(self, config):118        """Test loading data with invalid type"""119        config['data_source']['type'] = 'invalid_type'120        121        result = load_data(config)122        assert result is None123    124    def test_load_csv_data_file_exists(self, config, sample_csv_data):125        """Test loading CSV data when file exists"""126        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as tmp_file:127            sample_csv_data.to_csv(tmp_file.name, index=False)128            config['data_source']['path'] = tmp_file.name129            130            try:131                result = _load_csv_data(config)132                133                assert isinstance(result, pd.DataFrame)134                assert len(result) == len(sample_csv_data)135                assert pd.api.types.is_datetime64_any_dtype(result['timestamp'])136                137            finally:138                os.unlink(tmp_file.name)139    140    def test_load_csv_data_file_not_exists(self, config):141        """Test loading CSV data when file doesn't exist"""142        config['data_source']['path'] = 'nonexistent_file.csv'143        144        result = _load_csv_data(config)145        146        assert result is None147    148    def test_load_csv_data_missing_columns(self, config):149        """Test loading CSV data with missing columns"""150        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as tmp_file:151            # Create CSV with missing columns152            incomplete_data = pd.DataFrame({153                'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),154                'open': [150] * 10,155                'close': [152] * 10156                # Missing high, low, volume157            })158            incomplete_data.to_csv(tmp_file.name, index=False)159            config['data_source']['path'] = tmp_file.name160            161            try:162                result = _load_csv_data(config)163                164                assert result is None165                    166            finally:167                os.unlink(tmp_file.name)168    169    def test_load_synthetic_data(self, config):170        """Test synthetic data loading (mock generator and file existence)"""171        mock_df = pd.DataFrame({172            'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),173            'open': [150] * 10,174            'high': [155] * 10,175            'low': [145] * 10,176            'close': [152] * 10,177            'volume': [1000] * 10178        })179        with patch('os.path.exists', return_value=False):180            with patch('agentic_ai_system.synthetic_data_generator.SyntheticDataGenerator') as mock_generator_class:181                mock_generator = MagicMock()182                mock_generator_class.return_value = mock_generator183                mock_generator.generate_data.return_value = mock_df184 185                result = _load_synthetic_data(config)186                assert isinstance(result, pd.DataFrame)187                assert list(result.columns) == ['timestamp', 'open', 'high', 'low', 'close', 'volume']188    189    def test_validate_data_valid(self, sample_csv_data):190        """Test data validation with valid data"""191        # Create a copy to avoid modifying the original192        data_copy = sample_csv_data.copy()193        assert validate_data(data_copy) == True194    195    def test_validate_data_missing_columns(self):196        """Test data validation with missing columns"""197        invalid_data = pd.DataFrame({198            'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),199            'open': [150] * 10200            # Missing required columns201        })202        203        assert validate_data(invalid_data) == False204    205    def test_validate_data_negative_prices(self):206        """Test data validation with negative prices"""207        invalid_data = pd.DataFrame({208            'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),209            'open': [150] * 10,210            'high': [155] * 10,211            'low': [-145] * 10,  # Negative low price212            'close': [152] * 10,213            'volume': [1000] * 10214        })215        216        assert validate_data(invalid_data) == False217    218    def test_validate_data_negative_volumes(self):219        """Test data validation with negative volumes"""220        invalid_data = pd.DataFrame({221            'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),222            'open': [150] * 10,223            'high': [155] * 10,224            'low': [145] * 10,225            'close': [152] * 10,226            'volume': [-1000] * 10  # Negative volume227        })228        229        # The current implementation doesn't check for negative volumes230        # It only warns about high percentage of zero volumes231        assert validate_data(invalid_data) == True232    233    def test_validate_data_invalid_ohlc(self):234        """Test data validation with invalid OHLC relationships"""235        invalid_data = pd.DataFrame({236            'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),237            'open': [150] * 10,238            'high': [145] * 10,  # High < Open239            'low': [145] * 10,240            'close': [152] * 10,241            'volume': [1000] * 10242        })243        244        assert validate_data(invalid_data) == False245    246    def test_validate_data_null_values(self):247        """Test data validation with null values"""248        invalid_data = pd.DataFrame({249            'timestamp': pd.date_range('2024-01-01', periods=10, freq='1min'),250            'open': [150] * 10,251            'high': [155] * 10,252            'low': [145] * 10,253            'close': [152] * 10,254            'volume': [1000] * 10255        })256        257        # Add null values258        invalid_data.loc[0, 'open'] = None259        260        # The current implementation removes NaN values and continues261        # So it should return True after removing the NaN row262        result = validate_data(invalid_data)263        assert result == True264        # Check that the NaN row was removed265        assert len(invalid_data) == 9  # Original 10 - 1 NaN row266    267    def test_validate_data_empty_dataframe(self):268        """Test data validation with empty DataFrame"""269        empty_data = pd.DataFrame()270        assert validate_data(empty_data) == False271    272    def test_load_data_error_handling(self, config):273        """Test error handling in load_data"""274        config['data_source']['type'] = 'csv'275        config['data_source']['path'] = 'nonexistent_file.csv'276        277        result = load_data(config)278        assert result is None279    280    def test_csv_data_timestamp_conversion(self, config, sample_csv_data):281        """Test timestamp conversion in CSV loading"""282        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as tmp_file:283            # Convert timestamp to string for CSV284            sample_csv_data['timestamp'] = sample_csv_data['timestamp'].astype(str)285            sample_csv_data.to_csv(tmp_file.name, index=False)286            config['data_source']['path'] = tmp_file.name287            288            try:289                result = _load_csv_data(config)290                291                # Check that timestamp is converted to datetime292                assert pd.api.types.is_datetime64_any_dtype(result['timestamp'])293                294            finally:295                os.unlink(tmp_file.name)296    297    def test_synthetic_data_directory_creation(self, config):298        """Test that synthetic data directory is created if it doesn't exist"""299        with patch('os.makedirs') as mock_makedirs:300            with patch('agentic_ai_system.synthetic_data_generator.SyntheticDataGenerator') as mock_generator_class:301                mock_generator = MagicMock()302                mock_generator_class.return_value = mock_generator303                304                mock_df = pd.DataFrame({'test': [1, 2, 3]})305                mock_generator.generate_data.return_value = mock_df306                307                # Mock os.path.exists to return False so it generates new data308                with patch('os.path.exists', return_value=False):309                    _load_synthetic_data(config)310                    311                    # Check that makedirs was called312                    mock_makedirs.assert_called_once()313    314    def test_data_validation_edge_cases(self):315        """Test data validation with edge cases"""316        # Test with single row317        single_row_data = pd.DataFrame({318            'timestamp': [pd.Timestamp('2024-01-01')],319            'open': [150],320            'high': [155],321            'low': [145],322            'close': [152],323            'volume': [1000]324        })325        326        assert validate_data(single_row_data) == True327        328        # Test with very large numbers329        large_data = pd.DataFrame({330            'timestamp': pd.date_range('2024-01-01', periods=5, freq='1min'),331            'open': [1e6] * 5,332            'high': [1e6 + 100] * 5,333            'low': [1e6 - 100] * 5,334            'close': [1e6 + 50] * 5,335            'volume': [1e9] * 5336        })337        338        assert validate_data(large_data) == True