ParallelLLC/algorithmic_trading
2732
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 