Team Ai
Modelpublic

ParallelLLC/algorithmic_trading

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
27likes32downloads
test_strategy_agent.py252 linesDownload Raw Back to tests
1import pytest2import pandas as pd3import numpy as np4from datetime import datetime, timedelta5from agentic_ai_system.strategy_agent import StrategyAgent6 7class TestStrategyAgent:8    """Test cases for StrategyAgent"""9    10    @pytest.fixture11    def config(self):12        """Sample configuration for testing"""13        return {14            'trading': {15                'symbol': 'AAPL',16                'timeframe': '1min',17                'capital': 10000018            },19            'risk': {20                'max_position': 100,21                'max_drawdown': 0.0522            },23            'execution': {24                'broker_api': 'paper',25                'order_size': 1026            }27        }28    29    @pytest.fixture30    def strategy_agent(self, config):31        """Create a StrategyAgent instance"""32        return StrategyAgent(config)33    34    @pytest.fixture35    def sample_data(self):36        """Create sample market data for testing"""37        dates = pd.date_range(start='2024-01-01', periods=100, freq='1min')38        39        # Generate realistic price data40        base_price = 150.041        prices = []42        for i in range(100):43            # Add some trend and noise44            price = base_price + (i * 0.1) + np.random.normal(0, 2)45            prices.append(max(price, 1))  # Ensure positive prices46        47        data = []48        for i, (date, close_price) in enumerate(zip(dates, prices)):49            # Generate OHLC from close price50            noise = np.random.normal(0, 1)51            open_price = close_price + noise52            high_price = max(open_price, close_price) + abs(np.random.normal(0, 2))53            low_price = min(open_price, close_price) - abs(np.random.normal(0, 2))54            volume = np.random.randint(1000, 100000)55            56            data.append({57                'timestamp': date,58                'open': round(open_price, 2),59                'high': round(high_price, 2),60                'low': round(low_price, 2),61                'close': round(close_price, 2),62                'volume': volume63            })64        65        return pd.DataFrame(data)66    67    def test_initialization(self, strategy_agent, config):68        """Test agent initialization"""69        assert strategy_agent.symbol == config['trading']['symbol']70        assert strategy_agent.capital == config['trading']['capital']71        assert strategy_agent.max_position == config['risk']['max_position']72        assert strategy_agent.max_drawdown == config['risk']['max_drawdown']73    74    def test_act_with_valid_data(self, strategy_agent, sample_data):75        """Test signal generation with valid data"""76        signal = strategy_agent.act(sample_data)77        78        # Check signal structure79        assert isinstance(signal, dict)80        assert 'action' in signal81        assert 'symbol' in signal82        assert 'quantity' in signal83        assert 'price' in signal84        assert 'confidence' in signal85        86        # Check action values87        assert signal['action'] in ['buy', 'sell', 'hold']88        assert signal['symbol'] == strategy_agent.symbol89        assert signal['quantity'] >= 090        assert signal['price'] > 091        assert 0 <= signal['confidence'] <= 192    93    def test_act_with_empty_data(self, strategy_agent):94        """Test signal generation with empty data"""95        empty_data = pd.DataFrame()96        signal = strategy_agent.act(empty_data)97        98        assert signal['action'] == 'hold'99        assert signal['quantity'] == 0100        assert signal['confidence'] == 0.0101    102    def test_calculate_indicators(self, strategy_agent, sample_data):103        """Test technical indicator calculations"""104        indicators = strategy_agent._calculate_indicators(sample_data)105        106        # Check that indicators are calculated107        expected_indicators = ['sma_20', 'sma_50', 'rsi', 'bb_upper', 'bb_lower', 'macd', 'macd_signal']108        for indicator in expected_indicators:109            assert indicator in indicators110        111        # Check that indicators have reasonable values112        if len(indicators['sma_20']) > 0:113            assert indicators['sma_20'][-1] > 0114        115        if len(indicators['rsi']) > 0:116            rsi_value = indicators['rsi'][-1]117            assert 0 <= rsi_value <= 100118    119    def test_calculate_sma(self, strategy_agent):120        """Test Simple Moving Average calculation"""121        prices = np.array([100, 101, 102, 103, 104, 105, 106, 107, 108, 109])122        123        # Test SMA with window 3124        sma = strategy_agent._calculate_sma(prices, 3)125        expected_sma = np.array([101, 102, 103, 104, 105, 106, 107, 108])126        127        np.testing.assert_array_almost_equal(sma, expected_sma, decimal=2)128        129        # Test with insufficient data130        short_prices = np.array([100, 101])131        sma_short = strategy_agent._calculate_sma(short_prices, 3)132        assert len(sma_short) == 0133    134    def test_calculate_rsi(self, strategy_agent):135        """Test RSI calculation"""136        # Create price data with known pattern137        prices = np.array([100, 101, 102, 101, 100, 99, 98, 99, 100, 101])138        139        rsi = strategy_agent._calculate_rsi(prices, window=3)140        141        # RSI should be between 0 and 100142        if len(rsi) > 0:143            assert 0 <= rsi[-1] <= 100144    145    def test_calculate_bollinger_bands(self, strategy_agent):146        """Test Bollinger Bands calculation"""147        prices = np.array([100, 101, 102, 103, 104, 105, 106, 107, 108, 109])148        149        bb_upper, bb_lower = strategy_agent._calculate_bollinger_bands(prices, window=5)150        151        if len(bb_upper) > 0 and len(bb_lower) > 0:152            # Upper band should be above lower band153            assert bb_upper[-1] > bb_lower[-1]154    155    def test_calculate_position_size(self, strategy_agent):156        """Test position size calculation"""157        price = 150.0158        159        # Test normal case160        quantity = strategy_agent._calculate_position_size(price)161        expected_quantity = int((strategy_agent.capital * 0.1) / price)162        expected_quantity = min(expected_quantity, strategy_agent.max_position)163        164        assert quantity == expected_quantity165        assert quantity >= 1166        167        # Test with very high price168        high_price = 10000.0169        quantity_high = strategy_agent._calculate_position_size(high_price)170        assert quantity_high == 1  # Minimum quantity171    172    def test_generate_no_action_signal(self, strategy_agent):173        """Test no-action signal generation"""174        signal = strategy_agent._generate_no_action_signal()175        176        assert signal['action'] == 'hold'177        assert signal['quantity'] == 0178        assert signal['price'] == 0179        assert signal['confidence'] == 0.0180        assert signal['symbol'] == strategy_agent.symbol181    182    def test_signal_generation_logic(self, strategy_agent, sample_data):183        """Test signal generation logic with different market conditions"""184        # Test with upward trending data (should generate buy signal)185        upward_data = sample_data.copy()186        upward_data['close'] = upward_data['close'] * 1.1  # 10% increase187        188        signal_up = strategy_agent.act(upward_data)189        190        # Test with downward trending data (should generate sell signal)191        downward_data = sample_data.copy()192        downward_data['close'] = downward_data['close'] * 0.9  # 10% decrease193        194        signal_down = strategy_agent.act(downward_data)195        196        # Both should be valid signals197        assert signal_up['action'] in ['buy', 'sell', 'hold']198        assert signal_down['action'] in ['buy', 'sell', 'hold']199    200    def test_error_handling(self, strategy_agent):201        """Test error handling in signal generation"""202        # Test with invalid data203        invalid_data = pd.DataFrame({'invalid_column': [1, 2, 3]})204        205        # Should not raise exception, should return hold signal206        signal = strategy_agent.act(invalid_data)207        assert signal['action'] == 'hold'208    209    def test_technical_indicators_edge_cases(self, strategy_agent):210        """Test technical indicators with edge cases"""211        # Test with constant prices212        constant_prices = np.ones(50) * 100213        rsi_constant = strategy_agent._calculate_rsi(constant_prices)214        215        # Test with all increasing prices216        increasing_prices = np.arange(100, 150)217        rsi_increasing = strategy_agent._calculate_rsi(increasing_prices)218        219        # Test with all decreasing prices220        decreasing_prices = np.arange(150, 100, -1)221        rsi_decreasing = strategy_agent._calculate_rsi(decreasing_prices)222        223        # All should return valid arrays (possibly empty)224        assert isinstance(rsi_constant, np.ndarray)225        assert isinstance(rsi_increasing, np.ndarray)226        assert isinstance(rsi_decreasing, np.ndarray)227    228    def test_macd_calculation(self, strategy_agent):229        """Test MACD calculation"""230        prices = np.array([100 + i * 0.1 + np.random.normal(0, 1) for i in range(50)])231        232        macd, signal = strategy_agent._calculate_macd(prices)233        234        # Both should be numpy arrays235        assert isinstance(macd, np.ndarray)236        assert isinstance(signal, np.ndarray)237        238        # If we have enough data, both should have values239        if len(prices) >= 26:240            assert len(macd) > 0241            if len(macd) >= 9:  # Need enough data for signal line242                assert len(signal) > 0243    244    def test_ema_calculation(self, strategy_agent):245        """Test Exponential Moving Average calculation"""246        prices = np.array([100, 101, 102, 103, 104, 105, 106, 107, 108, 109])247        248        ema = strategy_agent._calculate_ema(prices, window=5)249        250        assert isinstance(ema, np.ndarray)251        if len(ema) > 0:252            assert ema[-1] > 0  # EMA should be positive