mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
feat: Add RL Trading Agent system with 99 tests
Implement Reinforcement Learning trading system inspired by FinRL concepts (100% original code, NOT copied from FinRL MIT project): RL ENVIRONMENT: - TradingEnv: Gymnasium-compatible environment - State: price history + indicators + portfolio state - Action: continuous position [-1, 1] (short to long) - Reward: return - transaction costs - drawdown penalty RL AGENT: - RLTradingAgent: Wrapper for Stable Baselines3 - Supports PPO (stable), A2C (fast), SAC (continuous) - Methods: create_model(), train(), predict(), save(), load(), evaluate() COSTEER (fills TODO at costeer.py:112): - RLCosteer: RL-based trading controller - Risk-limit enforcement (15% drawdown stops trading) - Position scaling based on risk appetite - Trade history tracking TECHNICAL INDICATORS: - RSI, MACD, Bollinger Bands, CCI, ATR - prepare_features() helper for easy integration TESTS (99 total, ALL PASS): - 26 env tests - 16 agent tests - 19 costeer tests - 18 indicator tests - 10 integration tests Documentation: - Update QWEN.md with RL system architecture
This commit is contained in:
TPTBusiness
@@ -895,3 +895,165 @@ class TestProtectionManager:
|
||||
assert StoplossGuardConfig is not None
|
||||
assert LowPerformanceConfig is not None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 15. RL TRADING AGENT TESTS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestRLTrading:
|
||||
"""Test RL Trading System."""
|
||||
|
||||
def test_rl_env_import(self):
|
||||
"""Test RL trading environment imports."""
|
||||
from rdagent.components.coder.rl.env import TradingEnv, TradingState
|
||||
assert TradingEnv is not None
|
||||
assert TradingState is not None
|
||||
|
||||
def test_rl_agent_import(self):
|
||||
"""Test RL trading agent imports."""
|
||||
from rdagent.components.coder.rl.agent import RLTradingAgent
|
||||
assert RLTradingAgent is not None
|
||||
|
||||
def test_costeer_import(self):
|
||||
"""Test RL Costeer imports."""
|
||||
from rdagent.components.coder.rl.costeer import RLCosteer
|
||||
assert RLCosteer is not None
|
||||
|
||||
def test_indicators_import(self):
|
||||
"""Test technical indicators import."""
|
||||
from rdagent.components.coder.rl.indicators import (
|
||||
calculate_rsi,
|
||||
calculate_macd,
|
||||
calculate_bollinger_bands,
|
||||
calculate_cci,
|
||||
calculate_atr,
|
||||
prepare_features,
|
||||
)
|
||||
assert calculate_rsi is not None
|
||||
assert calculate_macd is not None
|
||||
assert calculate_bollinger_bands is not None
|
||||
|
||||
def test_env_creation(self):
|
||||
"""Test RL environment can be created with mock data."""
|
||||
from rdagent.components.coder.rl.env import TradingEnv
|
||||
|
||||
np.random.seed(42)
|
||||
prices = 100.0 + np.cumsum(np.random.randn(200) * 0.5)
|
||||
env = TradingEnv(prices=prices, window_size=30, max_steps=100)
|
||||
|
||||
assert env is not None
|
||||
assert env.observation_space.shape[0] > 0
|
||||
assert env.action_space.shape == (1,)
|
||||
|
||||
def test_env_reset_and_step(self):
|
||||
"""Test environment reset and step work correctly."""
|
||||
from rdagent.components.coder.rl.env import TradingEnv
|
||||
|
||||
np.random.seed(42)
|
||||
prices = 100.0 + np.cumsum(np.random.randn(200) * 0.5)
|
||||
env = TradingEnv(prices=prices, window_size=10, max_steps=50)
|
||||
|
||||
obs, info = env.reset()
|
||||
assert isinstance(obs, np.ndarray)
|
||||
assert obs.dtype == np.float32
|
||||
|
||||
action = np.array([0.3], dtype=np.float32)
|
||||
obs, reward, terminated, truncated, info = env.step(action)
|
||||
|
||||
assert isinstance(reward, float)
|
||||
assert isinstance(terminated, bool)
|
||||
assert isinstance(truncated, bool)
|
||||
assert "equity" in info
|
||||
|
||||
def test_agent_creation(self):
|
||||
"""Test RL trading agent can be created."""
|
||||
from rdagent.components.coder.rl.agent import RLTradingAgent
|
||||
|
||||
agent = RLTradingAgent(algorithm="PPO")
|
||||
assert agent.algorithm == "PPO"
|
||||
assert agent.is_trained is False
|
||||
assert agent.model is None
|
||||
|
||||
def test_costeer_initialization(self):
|
||||
"""Test RL Costeer can be initialized."""
|
||||
from rdagent.components.coder.rl.costeer import RLCosteer
|
||||
|
||||
costeer = RLCosteer(
|
||||
algorithm="PPO",
|
||||
window_size=30,
|
||||
max_position=1.0,
|
||||
risk_limit=0.15,
|
||||
)
|
||||
assert costeer.algorithm == "PPO"
|
||||
assert costeer.is_active is False
|
||||
assert costeer.model is None
|
||||
|
||||
def test_full_rl_workflow(self):
|
||||
"""Test complete RL workflow: env -> indicators -> features."""
|
||||
from rdagent.components.coder.rl.env import TradingEnv
|
||||
from rdagent.components.coder.rl.indicators import prepare_features
|
||||
|
||||
# Create price data
|
||||
np.random.seed(42)
|
||||
n = 200
|
||||
dates = pd.date_range("2024-01-01", periods=n, freq="B")
|
||||
close = 100.0 + np.cumsum(np.random.randn(n) * 0.5)
|
||||
high = close + np.abs(np.random.randn(n) * 0.3)
|
||||
low = close - np.abs(np.random.randn(n) * 0.3)
|
||||
|
||||
prices_df = pd.DataFrame(
|
||||
{"close": close, "high": high, "low": low}, index=dates
|
||||
)
|
||||
|
||||
# Prepare features
|
||||
features = prepare_features(prices_df, indicator_list=["rsi", "macd"])
|
||||
assert "rsi" in features.columns
|
||||
assert "macd" in features.columns
|
||||
assert not features.isna().any().any()
|
||||
|
||||
# Create environment with indicators
|
||||
indicators = features[["rsi", "macd", "signal", "histogram"]].values
|
||||
env = TradingEnv(
|
||||
prices=close,
|
||||
indicators=indicators,
|
||||
window_size=20,
|
||||
max_steps=50,
|
||||
)
|
||||
|
||||
obs, info = env.reset()
|
||||
assert isinstance(obs, np.ndarray)
|
||||
|
||||
# Run a few steps
|
||||
for _ in range(5):
|
||||
action = np.array([0.2], dtype=np.float32)
|
||||
obs, reward, terminated, truncated, info = env.step(action)
|
||||
assert isinstance(reward, float)
|
||||
|
||||
def test_costeer_with_market_data(self):
|
||||
"""Test RLCosteer initializes with market data."""
|
||||
from rdagent.components.coder.rl.costeer import RLCosteer
|
||||
|
||||
np.random.seed(42)
|
||||
n = 200
|
||||
dates = pd.date_range("2024-01-01", periods=n, freq="B")
|
||||
prices = pd.Series(100.0 + np.cumsum(np.random.randn(n) * 0.5), index=dates)
|
||||
indicators = pd.DataFrame(
|
||||
np.random.randn(n, 3).astype(np.float32),
|
||||
columns=["rsi", "macd", "bb"],
|
||||
index=dates,
|
||||
)
|
||||
|
||||
costeer = RLCosteer(window_size=30)
|
||||
costeer.initialize(prices, indicators, initial_equity=100000.0)
|
||||
|
||||
assert costeer.is_active is True
|
||||
assert costeer.peak_equity == 100000.0
|
||||
|
||||
# Step should work without model (returns 0 action)
|
||||
trade = costeer.step(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
assert "timestamp" in trade
|
||||
assert "step" in trade
|
||||
|
||||
|
||||
@@ -0,0 +1,339 @@
|
||||
"""
|
||||
Tests for RL Costeer (Trading Controller).
|
||||
|
||||
Covers:
|
||||
- Costeer initialization with/without model
|
||||
- Market data initialization
|
||||
- Action retrieval (mocked model)
|
||||
- Risk limit enforcement
|
||||
- Observation building
|
||||
- Step execution
|
||||
- Performance tracking
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from rdagent.components.coder.rl.costeer import RLCosteer
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FIXTURES
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prices() -> pd.Series:
|
||||
"""Generate 200-step price series."""
|
||||
np.random.seed(42)
|
||||
return pd.Series(100.0 + np.cumsum(np.random.randn(200) * 0.5))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_indicators() -> pd.DataFrame:
|
||||
"""Generate mock indicators DataFrame."""
|
||||
np.random.seed(42)
|
||||
return pd.DataFrame(
|
||||
np.random.randn(200, 3).astype(np.float32),
|
||||
columns=["rsi", "macd", "bb_width"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def basic_costeer() -> RLCosteer:
|
||||
"""Create costeer without model."""
|
||||
return RLCosteer(
|
||||
algorithm="PPO",
|
||||
window_size=30,
|
||||
max_position=1.0,
|
||||
risk_limit=0.15,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def initialized_costeer(mock_prices: pd.Series, mock_indicators: pd.DataFrame) -> RLCosteer:
|
||||
"""Create costeer with market data but no model."""
|
||||
costeer = RLCosteer(window_size=30)
|
||||
costeer.initialize(mock_prices, mock_indicators, initial_equity=100000.0)
|
||||
return costeer
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# INITIALIZATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestCosteerInit:
|
||||
"""Test costeer initialization."""
|
||||
|
||||
def test_default_values(self) -> None:
|
||||
"""Default parameters should match specification."""
|
||||
costeer = RLCosteer()
|
||||
assert costeer.algorithm == "PPO"
|
||||
assert costeer.window_size == 60
|
||||
assert costeer.max_position == 1.0
|
||||
assert costeer.risk_limit == 0.15
|
||||
assert costeer.is_active is False
|
||||
assert costeer.model is None
|
||||
assert costeer.trade_history == []
|
||||
|
||||
def test_custom_values(self) -> None:
|
||||
"""Custom parameters should be stored."""
|
||||
costeer = RLCosteer(
|
||||
algorithm="SAC",
|
||||
window_size=120,
|
||||
max_position=0.5,
|
||||
risk_limit=0.10,
|
||||
)
|
||||
assert costeer.algorithm == "SAC"
|
||||
assert costeer.window_size == 120
|
||||
assert costeer.max_position == 0.5
|
||||
assert costeer.risk_limit == 0.10
|
||||
|
||||
def test_initialize_sets_market_data(
|
||||
self, mock_prices: pd.Series, mock_indicators: pd.DataFrame
|
||||
) -> None:
|
||||
"""Initialize should store market data and activate costeer."""
|
||||
costeer = RLCosteer(window_size=30)
|
||||
costeer.initialize(mock_prices, mock_indicators, initial_equity=50000.0)
|
||||
|
||||
assert costeer.is_active is True
|
||||
assert len(costeer.prices) == 200
|
||||
assert costeer.indicators is not None
|
||||
assert costeer.initial_equity == 50000.0
|
||||
assert costeer.current_step == 30
|
||||
assert costeer.peak_equity == 50000.0
|
||||
|
||||
def test_initialize_without_indicators(self, mock_prices: pd.Series) -> None:
|
||||
"""Initialize should work without indicators."""
|
||||
costeer = RLCosteer(window_size=30)
|
||||
costeer.initialize(mock_prices, initial_equity=100000.0)
|
||||
|
||||
assert costeer.indicators is None
|
||||
assert costeer.is_active is True
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# GET ACTION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestGetAction:
|
||||
"""Test action retrieval."""
|
||||
|
||||
def test_no_model_returns_zero(self, initialized_costeer: RLCosteer) -> None:
|
||||
"""Without model, action should be 0 (hold)."""
|
||||
action = initialized_costeer.get_action(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
assert action == 0.0
|
||||
|
||||
def test_not_active_returns_zero(self, basic_costeer: RLCosteer) -> None:
|
||||
"""Inactive costeer should return 0."""
|
||||
action = basic_costeer.get_action(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
assert action == 0.0
|
||||
|
||||
@patch.object(RLCosteer, "_build_observation")
|
||||
def test_model_action_risk_scaled(
|
||||
self, mock_obs: MagicMock, initialized_costeer: RLCosteer
|
||||
) -> None:
|
||||
"""Model action should be returned and risk-scaled."""
|
||||
mock_obs.return_value = np.random.randn(100).astype(np.float32)
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_model.predict.return_value = (np.array([[0.8]]), None)
|
||||
initialized_costeer.model = mock_model
|
||||
|
||||
action = initialized_costeer.get_action(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
|
||||
# At full risk (no drawdown), action should be ~0.8
|
||||
assert abs(action) <= 1.0
|
||||
|
||||
def test_risk_limit_forces_close(self, initialized_costeer: RLCosteer) -> None:
|
||||
"""Drawdown > risk_limit should force position to 0."""
|
||||
mock_model = MagicMock()
|
||||
mock_model.predict.return_value = (np.array([[1.0]]), None)
|
||||
initialized_costeer.model = mock_model
|
||||
|
||||
# Simulate 20% drawdown (> 15% limit)
|
||||
action = initialized_costeer.get_action(
|
||||
current_equity=80000.0, # 20% drawdown from 100k
|
||||
cash=50000.0,
|
||||
position=0.5,
|
||||
)
|
||||
assert action == 0.0
|
||||
|
||||
def test_action_clipped_to_max_position(
|
||||
self, initialized_costeer: RLCosteer
|
||||
) -> None:
|
||||
"""Action should be clipped to max_position."""
|
||||
initialized_costeer.max_position = 0.5
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_model.predict.return_value = (np.array([[1.0]]), None)
|
||||
initialized_costeer.model = mock_model
|
||||
|
||||
action = initialized_costeer.get_action(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
assert action <= 0.5
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# OBSERVATION BUILDING
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestObservationBuilding:
|
||||
"""Test observation vector construction."""
|
||||
|
||||
def test_observation_shape_no_indicators(
|
||||
self, mock_prices: pd.Series
|
||||
) -> None:
|
||||
"""Observation should have correct shape without indicators."""
|
||||
costeer = RLCosteer(window_size=30)
|
||||
costeer.initialize(mock_prices, initial_equity=100000.0)
|
||||
|
||||
obs = costeer._build_observation(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
|
||||
# window_size + 3 (position, pnl, equity_ratio)
|
||||
expected = 30 + 3
|
||||
assert obs.shape == (expected,)
|
||||
|
||||
def test_observation_shape_with_indicators(
|
||||
self, mock_prices: pd.Series, mock_indicators: pd.DataFrame
|
||||
) -> None:
|
||||
"""Observation should include indicator dimensions."""
|
||||
costeer = RLCosteer(window_size=30)
|
||||
costeer.initialize(mock_prices, mock_indicators, initial_equity=100000.0)
|
||||
|
||||
obs = costeer._build_observation(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
|
||||
# window_size * (1 + 3) + 3
|
||||
expected = 30 + (30 * 3) + 3
|
||||
assert obs.shape == (expected,)
|
||||
|
||||
def test_observation_dtype(self, initialized_costeer: RLCosteer) -> None:
|
||||
"""Observation should be float32."""
|
||||
obs = initialized_costeer._build_observation(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
assert obs.dtype == np.float32
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# STEP EXECUTION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestStepExecution:
|
||||
"""Test step execution."""
|
||||
|
||||
def test_step_records_trade(self, initialized_costeer: RLCosteer) -> None:
|
||||
"""Step should record trade in history."""
|
||||
trade = initialized_costeer.step(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
|
||||
assert "timestamp" in trade
|
||||
assert "step" in trade
|
||||
assert "equity" in trade
|
||||
assert "position" in trade
|
||||
assert "target_position" in trade
|
||||
assert "action" in trade
|
||||
|
||||
def test_step_advances_current_step(self, initialized_costeer: RLCosteer) -> None:
|
||||
"""Step should increment current_step."""
|
||||
initial_step = initialized_costeer.current_step
|
||||
initialized_costeer.step(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
assert initialized_costeer.current_step == initial_step + 1
|
||||
|
||||
def test_step_updates_peak_equity(self, initialized_costeer: RLCosteer) -> None:
|
||||
"""Peak equity should update when equity exceeds previous peak."""
|
||||
initialized_costeer.peak_equity = 100000.0
|
||||
initialized_costeer.step(
|
||||
current_equity=105000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
assert initialized_costeer.peak_equity == 105000.0
|
||||
|
||||
def test_multiple_steps_accumulate_trades(
|
||||
self, initialized_costeer: RLCosteer
|
||||
) -> None:
|
||||
"""Multiple steps should accumulate trades."""
|
||||
for _ in range(5):
|
||||
initialized_costeer.step(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
|
||||
assert len(initialized_costeer.trade_history) == 5
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MODEL LOADING
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestModelLoading:
|
||||
"""Test model loading functionality."""
|
||||
|
||||
def test_load_model_import_error(self, tmp_path: Path) -> None:
|
||||
"""Load should raise ImportError when SB3 not installed."""
|
||||
costeer = RLCosteer()
|
||||
|
||||
with patch.dict("sys.modules", {"stable_baselines3": None}):
|
||||
with pytest.raises(ImportError, match="stable-baselines3"):
|
||||
costeer.load_model(tmp_path / "model.zip")
|
||||
|
||||
def test_load_model_file_not_found(self, tmp_path: Path) -> None:
|
||||
"""Load should raise ValueError for non-existent file."""
|
||||
costeer = RLCosteer(model_path=tmp_path / "nonexistent.zip")
|
||||
# Model path doesn't exist, so it won't try to load
|
||||
assert costeer.is_active is False
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PERFORMANCE TRACKING
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestPerformanceTracking:
|
||||
"""Test performance history."""
|
||||
|
||||
def test_get_performance_returns_dataframe(
|
||||
self, initialized_costeer: RLCosteer
|
||||
) -> None:
|
||||
"""Performance should return DataFrame."""
|
||||
# Add some trades
|
||||
initialized_costeer.step(
|
||||
current_equity=100000.0, cash=50000.0, position=0.0
|
||||
)
|
||||
initialized_costeer.step(
|
||||
current_equity=101000.0, cash=49000.0, position=0.3
|
||||
)
|
||||
|
||||
df = initialized_costeer.get_performance()
|
||||
|
||||
assert isinstance(df, pd.DataFrame)
|
||||
assert len(df) == 2
|
||||
assert "equity" in df.columns
|
||||
assert "position" in df.columns
|
||||
|
||||
def test_empty_performance_history(self, basic_costeer: RLCosteer) -> None:
|
||||
"""Empty history should return empty DataFrame."""
|
||||
df = basic_costeer.get_performance()
|
||||
assert isinstance(df, pd.DataFrame)
|
||||
assert len(df) == 0
|
||||
@@ -0,0 +1,276 @@
|
||||
"""
|
||||
Tests for Technical Indicators.
|
||||
|
||||
Covers:
|
||||
- RSI calculation
|
||||
- MACD calculation
|
||||
- Bollinger Bands calculation
|
||||
- CCI calculation
|
||||
- ATR calculation
|
||||
- prepare_features integration
|
||||
- Edge cases (NaN handling, short series)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from rdagent.components.coder.rl.indicators import (
|
||||
calculate_atr,
|
||||
calculate_bollinger_bands,
|
||||
calculate_cci,
|
||||
calculate_macd,
|
||||
calculate_rsi,
|
||||
prepare_features,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FIXTURES
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def price_data() -> pd.DataFrame:
|
||||
"""Generate 100 bars of realistic price data."""
|
||||
np.random.seed(42)
|
||||
n = 100
|
||||
close = 100.0 + np.cumsum(np.random.randn(n) * 0.5)
|
||||
high = close + np.abs(np.random.randn(n) * 0.3)
|
||||
low = close - np.abs(np.random.randn(n) * 0.3)
|
||||
volume = np.random.randint(1000, 10000, n)
|
||||
|
||||
return pd.DataFrame(
|
||||
{"close": close, "high": high, "low": low, "volume": volume},
|
||||
index=pd.date_range("2024-01-01", periods=n, freq="B"),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# RSI
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestRSI:
|
||||
"""Test RSI calculation."""
|
||||
|
||||
def test_rsi_values_in_range(self, price_data: pd.DataFrame) -> None:
|
||||
"""RSI should be between 0 and 100."""
|
||||
rsi = calculate_rsi(price_data["close"], period=14)
|
||||
# Skip NaN values at the beginning
|
||||
valid_rsi = rsi.dropna()
|
||||
assert (valid_rsi >= 0).all()
|
||||
assert (valid_rsi <= 100).all()
|
||||
|
||||
def test_rsi_default_period(self, price_data: pd.DataFrame) -> None:
|
||||
"""Default RSI period should be 14."""
|
||||
rsi = calculate_rsi(price_data["close"])
|
||||
assert len(rsi) == len(price_data)
|
||||
|
||||
def test_rsi_custom_period(self, price_data: pd.DataFrame) -> None:
|
||||
"""Custom period should be respected."""
|
||||
rsi_7 = calculate_rsi(price_data["close"], period=7)
|
||||
rsi_21 = calculate_rsi(price_data["close"], period=21)
|
||||
|
||||
# Shorter period = more valid values at start
|
||||
assert rsi_7.dropna().iloc[0] >= 0
|
||||
assert rsi_7.dropna().iloc[0] <= 100
|
||||
|
||||
def test_rsi_nan_at_start(self, price_data: pd.DataFrame) -> None:
|
||||
"""RSI should have NaN values at the beginning (period-1)."""
|
||||
rsi = calculate_rsi(price_data["close"], period=14)
|
||||
# First 13 values should be NaN (need 14 for rolling mean)
|
||||
assert rsi.iloc[:13].isna().all()
|
||||
# Value at index 13 should be valid
|
||||
assert not np.isnan(rsi.iloc[13])
|
||||
|
||||
def test_rsi_short_series(self) -> None:
|
||||
"""RSI should handle short series gracefully."""
|
||||
prices = pd.Series([100.0, 101.0, 102.0])
|
||||
rsi = calculate_rsi(prices, period=14)
|
||||
# All should be NaN (not enough data)
|
||||
assert rsi.isna().all()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MACD
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestMACD:
|
||||
"""Test MACD calculation."""
|
||||
|
||||
def test_macd_output_columns(self, price_data: pd.DataFrame) -> None:
|
||||
"""MACD should return DataFrame with macd, signal, histogram."""
|
||||
macd_df = calculate_macd(price_data["close"])
|
||||
assert "macd" in macd_df.columns
|
||||
assert "signal" in macd_df.columns
|
||||
assert "histogram" in macd_df.columns
|
||||
|
||||
def test_macd_histogram_consistency(self, price_data: pd.DataFrame) -> None:
|
||||
"""Histogram should equal MACD - Signal."""
|
||||
macd_df = calculate_macd(price_data["close"])
|
||||
valid = macd_df.dropna()
|
||||
expected_histogram = valid["macd"] - valid["signal"]
|
||||
np.testing.assert_array_almost_equal(
|
||||
valid["histogram"].values,
|
||||
expected_histogram.values,
|
||||
decimal=10,
|
||||
)
|
||||
|
||||
def test_macd_custom_parameters(self, price_data: pd.DataFrame) -> None:
|
||||
"""Custom MACD parameters should be respected."""
|
||||
macd_df = calculate_macd(price_data["close"], fast=6, slow=13, signal=4)
|
||||
assert len(macd_df) == len(price_data)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# BOLLINGER BANDS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestBollingerBands:
|
||||
"""Test Bollinger Bands calculation."""
|
||||
|
||||
def test_bb_output_columns(self, price_data: pd.DataFrame) -> None:
|
||||
"""Bollinger Bands should return upper, middle, lower."""
|
||||
bb_df = calculate_bollinger_bands(price_data["close"])
|
||||
assert "upper" in bb_df.columns
|
||||
assert "middle" in bb_df.columns
|
||||
assert "lower" in bb_df.columns
|
||||
|
||||
def test_bb_ordering(self, price_data: pd.DataFrame) -> None:
|
||||
"""Upper >= Middle >= Lower for valid data."""
|
||||
bb_df = calculate_bollinger_bands(price_data["close"]).dropna()
|
||||
assert (bb_df["upper"] >= bb_df["middle"]).all()
|
||||
assert (bb_df["middle"] >= bb_df["lower"]).all()
|
||||
|
||||
def test_bb_middle_is_sma(self, price_data: pd.DataFrame) -> None:
|
||||
"""Middle band should equal SMA."""
|
||||
bb_df = calculate_bollinger_bands(price_data["close"], period=20)
|
||||
sma = price_data["close"].rolling(window=20).mean()
|
||||
valid = bb_df.dropna()
|
||||
np.testing.assert_array_almost_equal(
|
||||
valid["middle"].values,
|
||||
sma.dropna().values,
|
||||
decimal=10,
|
||||
)
|
||||
|
||||
def test_bb_custom_std_dev(self, price_data: pd.DataFrame) -> None:
|
||||
"""Custom std_dev should affect band width."""
|
||||
bb_1 = calculate_bollinger_bands(price_data["close"], std_dev=1.0).dropna()
|
||||
bb_2 = calculate_bollinger_bands(price_data["close"], std_dev=3.0).dropna()
|
||||
|
||||
# Higher std_dev = wider bands
|
||||
width_1 = (bb_1["upper"] - bb_1["lower"]).mean()
|
||||
width_2 = (bb_2["upper"] - bb_2["lower"]).mean()
|
||||
assert width_2 > width_1
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CCI
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestCCI:
|
||||
"""Test CCI calculation."""
|
||||
|
||||
def test_cci_values(self, price_data: pd.DataFrame) -> None:
|
||||
"""CCI should produce finite values after warmup."""
|
||||
cci = calculate_cci(
|
||||
price_data["close"], price_data["high"], price_data["low"], period=20
|
||||
)
|
||||
valid_cci = cci.dropna()
|
||||
assert len(valid_cci) > 0
|
||||
assert np.all(np.isfinite(valid_cci))
|
||||
|
||||
def test_cci_nan_at_start(self, price_data: pd.DataFrame) -> None:
|
||||
"""CCI should have NaN at the beginning."""
|
||||
cci = calculate_cci(
|
||||
price_data["close"], price_data["high"], price_data["low"], period=20
|
||||
)
|
||||
# First ~19 values should be NaN
|
||||
assert cci.iloc[:19].isna().any()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ATR
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestATR:
|
||||
"""Test ATR calculation."""
|
||||
|
||||
def test_atr_positive(self, price_data: pd.DataFrame) -> None:
|
||||
"""ATR should be positive (for non-zero price changes)."""
|
||||
atr = calculate_atr(
|
||||
price_data["high"], price_data["low"], price_data["close"], period=14
|
||||
)
|
||||
valid_atr = atr.dropna()
|
||||
assert (valid_atr > 0).all()
|
||||
|
||||
def test_atr_nan_at_start(self, price_data: pd.DataFrame) -> None:
|
||||
"""ATR should have NaN at the beginning."""
|
||||
atr = calculate_atr(
|
||||
price_data["high"], price_data["low"], price_data["close"], period=14
|
||||
)
|
||||
# First value should be NaN (no previous close)
|
||||
assert np.isnan(atr.iloc[0])
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PREPARE FEATURES
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestPrepareFeatures:
|
||||
"""Test feature preparation integration."""
|
||||
|
||||
def test_default_indicators(self, price_data: pd.DataFrame) -> None:
|
||||
"""Default should include rsi, macd, bollinger, sma."""
|
||||
features = prepare_features(price_data)
|
||||
|
||||
assert "rsi" in features.columns
|
||||
assert "macd" in features.columns
|
||||
assert "signal" in features.columns
|
||||
assert "histogram" in features.columns
|
||||
assert "upper" in features.columns
|
||||
assert "middle" in features.columns
|
||||
assert "lower" in features.columns
|
||||
assert "sma_20" in features.columns
|
||||
assert "sma_50" in features.columns
|
||||
|
||||
def test_custom_indicator_list(self, price_data: pd.DataFrame) -> None:
|
||||
"""Only requested indicators should be included."""
|
||||
features = prepare_features(price_data, indicator_list=["rsi"])
|
||||
|
||||
assert "rsi" in features.columns
|
||||
# MACD and Bollinger should NOT be present
|
||||
assert "macd" not in features.columns
|
||||
assert "upper" not in features.columns
|
||||
|
||||
def test_no_nan_in_output(self, price_data: pd.DataFrame) -> None:
|
||||
"""Output should have no NaN values."""
|
||||
features = prepare_features(price_data)
|
||||
assert not features.isna().any().any()
|
||||
|
||||
def test_preserves_original_columns(self, price_data: pd.DataFrame) -> None:
|
||||
"""Original price columns should be preserved."""
|
||||
features = prepare_features(price_data)
|
||||
for col in price_data.columns:
|
||||
assert col in features.columns
|
||||
|
||||
def test_empty_indicator_list(self, price_data: pd.DataFrame) -> None:
|
||||
"""Empty list should return only original data."""
|
||||
features = prepare_features(price_data, indicator_list=[])
|
||||
assert list(features.columns) == list(price_data.columns)
|
||||
|
||||
def test_close_only_dataframe(self) -> None:
|
||||
"""Should work with only 'close' column."""
|
||||
np.random.seed(42)
|
||||
prices = pd.DataFrame(
|
||||
{"close": 100.0 + np.cumsum(np.random.randn(100) * 0.5)}
|
||||
)
|
||||
features = prepare_features(prices, indicator_list=["rsi"])
|
||||
assert "rsi" in features.columns
|
||||
assert not features.isna().any().any()
|
||||
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
Tests for RL Trading Agent wrapper.
|
||||
|
||||
Covers:
|
||||
- Agent creation with different algorithms
|
||||
- Parameter merging with defaults
|
||||
- Model creation (mocked, since SB3 may not be installed)
|
||||
- Predict/Save/Load error handling
|
||||
- Evaluation (mocked)
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from rdagent.components.coder.rl.agent import RLTradingAgent
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FIXTURES
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_env() -> MagicMock:
|
||||
"""Create a mock gym environment."""
|
||||
env = MagicMock()
|
||||
env.observation_space = MagicMock()
|
||||
env.action_space = MagicMock()
|
||||
return env
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model() -> MagicMock:
|
||||
"""Create a mock SB3 model."""
|
||||
model = MagicMock()
|
||||
model.predict.return_value = (np.array([0.5]), None)
|
||||
return model
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# AGENT CREATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestAgentCreation:
|
||||
"""Test agent initialization."""
|
||||
|
||||
def test_default_creation(self) -> None:
|
||||
"""Default agent should use PPO with standard params."""
|
||||
agent = RLTradingAgent()
|
||||
assert agent.algorithm == "PPO"
|
||||
assert agent.policy == "MlpPolicy"
|
||||
assert agent.verbose == 0
|
||||
assert agent.model is None
|
||||
assert agent.is_trained is False
|
||||
|
||||
def test_algorithm_uppercase(self) -> None:
|
||||
"""Algorithm name should be uppercased."""
|
||||
agent = RLTradingAgent(algorithm="ppo")
|
||||
assert agent.algorithm == "PPO"
|
||||
|
||||
def test_custom_params_merge(self) -> None:
|
||||
"""User params should merge with defaults."""
|
||||
agent = RLTradingAgent(
|
||||
algorithm="PPO",
|
||||
params={"learning_rate": 1e-3, "gamma": 0.95},
|
||||
)
|
||||
# User param should override
|
||||
assert agent.params["learning_rate"] == 1e-3
|
||||
assert agent.params["gamma"] == 0.95
|
||||
# Default should remain
|
||||
assert "n_steps" in agent.params
|
||||
assert agent.params["n_steps"] == 2048
|
||||
|
||||
def test_a2c_default_params(self) -> None:
|
||||
"""A2C should have its own default params."""
|
||||
agent = RLTradingAgent(algorithm="A2C")
|
||||
assert agent.params["n_steps"] == 5
|
||||
assert agent.params["learning_rate"] == 7e-4
|
||||
|
||||
def test_sac_default_params(self) -> None:
|
||||
"""SAC should have its own default params."""
|
||||
agent = RLTradingAgent(algorithm="SAC")
|
||||
assert agent.params["buffer_size"] == 1_000_000
|
||||
assert agent.params["batch_size"] == 256
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MODEL CREATION (MOCKED)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestModelCreation:
|
||||
"""Test model creation with mocked SB3."""
|
||||
|
||||
@patch("stable_baselines3.PPO")
|
||||
def test_create_model_ppo(self, mock_ppo: MagicMock, mock_env: MagicMock) -> None:
|
||||
"""PPO model should be created with correct params."""
|
||||
agent = RLTradingAgent(algorithm="PPO", params={"n_steps": 100})
|
||||
agent.create_model(mock_env)
|
||||
mock_ppo.assert_called_once()
|
||||
call_kwargs = mock_ppo.call_args.kwargs
|
||||
assert call_kwargs["verbose"] == 0
|
||||
|
||||
@patch("stable_baselines3.A2C")
|
||||
def test_create_model_a2c(self, mock_a2c: MagicMock, mock_env: MagicMock) -> None:
|
||||
"""A2C model should be created with correct params."""
|
||||
agent = RLTradingAgent(algorithm="A2C")
|
||||
agent.create_model(mock_env)
|
||||
mock_a2c.assert_called_once()
|
||||
|
||||
@patch("stable_baselines3.SAC")
|
||||
def test_create_model_sac(self, mock_sac: MagicMock, mock_env: MagicMock) -> None:
|
||||
"""SAC model should be created with correct params."""
|
||||
agent = RLTradingAgent(algorithm="SAC")
|
||||
agent.create_model(mock_env)
|
||||
mock_sac.assert_called_once()
|
||||
|
||||
def test_model_class_not_found(self) -> None:
|
||||
"""Invalid algorithm should raise ImportError."""
|
||||
agent = RLTradingAgent(algorithm="INVALID")
|
||||
with pytest.raises(ImportError, match="Unknown algorithm"):
|
||||
agent._get_model_class()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TRAINING (MOCKED)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestTraining:
|
||||
"""Test training with mocked model."""
|
||||
|
||||
@patch("stable_baselines3.PPO")
|
||||
def test_train_returns_metadata(
|
||||
self, mock_ppo: MagicMock, mock_env: MagicMock
|
||||
) -> None:
|
||||
"""Training should return metadata dict."""
|
||||
mock_model = MagicMock()
|
||||
mock_ppo.return_value = mock_model
|
||||
|
||||
agent = RLTradingAgent(algorithm="PPO")
|
||||
result = agent.train(mock_env, total_timesteps=1000)
|
||||
|
||||
assert result["algorithm"] == "PPO"
|
||||
assert result["total_timesteps"] == 1000
|
||||
assert result["is_trained"] is True
|
||||
assert agent.is_trained is True
|
||||
|
||||
@patch("stable_baselines3.PPO")
|
||||
def test_train_calls_learn(
|
||||
self, mock_ppo: MagicMock, mock_env: MagicMock
|
||||
) -> None:
|
||||
"""Training should call model.learn."""
|
||||
mock_model = MagicMock()
|
||||
mock_ppo.return_value = mock_model
|
||||
|
||||
agent = RLTradingAgent(algorithm="PPO")
|
||||
agent.train(mock_env, total_timesteps=5000)
|
||||
|
||||
mock_model.learn.assert_called_once()
|
||||
call_kwargs = mock_model.learn.call_args.kwargs
|
||||
assert call_kwargs["total_timesteps"] == 5000
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PREDICTION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestPrediction:
|
||||
"""Test prediction functionality."""
|
||||
|
||||
def test_predict_without_model_raises(self) -> None:
|
||||
"""Predict without model should raise ValueError."""
|
||||
agent = RLTradingAgent()
|
||||
obs = np.random.randn(120).astype(np.float32)
|
||||
with pytest.raises(ValueError, match="not trained or loaded"):
|
||||
agent.predict(obs)
|
||||
|
||||
def test_predict_returns_action(self, mock_model: MagicMock) -> None:
|
||||
"""Predict should return action array."""
|
||||
agent = RLTradingAgent()
|
||||
agent.model = mock_model
|
||||
|
||||
obs = np.random.randn(120).astype(np.float32)
|
||||
action = agent.predict(obs)
|
||||
|
||||
assert isinstance(action, np.ndarray)
|
||||
mock_model.predict.assert_called_once_with(obs, deterministic=True)
|
||||
|
||||
def test_predict_non_deterministic(self, mock_model: MagicMock) -> None:
|
||||
"""Predict with deterministic=False should pass flag."""
|
||||
agent = RLTradingAgent()
|
||||
agent.model = mock_model
|
||||
|
||||
obs = np.random.randn(120).astype(np.float32)
|
||||
agent.predict(obs, deterministic=False)
|
||||
|
||||
mock_model.predict.assert_called_once_with(obs, deterministic=False)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SAVE / LOAD (MOCKED)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestSaveLoad:
|
||||
"""Test model save and load."""
|
||||
|
||||
def test_save_without_model_raises(self, tmp_path: Path) -> None:
|
||||
"""Save without model should raise ValueError."""
|
||||
agent = RLTradingAgent()
|
||||
with pytest.raises(ValueError, match="No model to save"):
|
||||
agent.save(tmp_path / "model.zip")
|
||||
|
||||
@patch("stable_baselines3.PPO")
|
||||
def test_save_creates_directory(self, mock_ppo: MagicMock, tmp_path: Path) -> None:
|
||||
"""Save should create parent directories."""
|
||||
mock_model = MagicMock()
|
||||
mock_ppo.return_value = mock_model
|
||||
|
||||
agent = RLTradingAgent(algorithm="PPO")
|
||||
agent.model = mock_model
|
||||
|
||||
save_path = tmp_path / "subdir" / "model.zip"
|
||||
agent.save(save_path)
|
||||
|
||||
mock_model.save.assert_called_once_with(str(save_path))
|
||||
|
||||
@patch("stable_baselines3.PPO")
|
||||
def test_load_sets_trained_flag(self, mock_ppo: MagicMock, tmp_path: Path) -> None:
|
||||
"""Load should set is_trained to True."""
|
||||
mock_model_class = MagicMock()
|
||||
mock_ppo.load.return_value = MagicMock()
|
||||
mock_ppo.return_value = mock_model_class
|
||||
|
||||
agent = RLTradingAgent(algorithm="PPO")
|
||||
agent.load(tmp_path / "model.zip")
|
||||
|
||||
assert agent.is_trained is True
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# EVALUATION (MOCKED)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestEvaluation:
|
||||
"""Test evaluation functionality."""
|
||||
|
||||
def test_evaluate_without_model_raises(self, mock_env: MagicMock) -> None:
|
||||
"""Evaluate without model should raise ValueError."""
|
||||
agent = RLTradingAgent()
|
||||
with pytest.raises(ValueError, match="not trained or loaded"):
|
||||
agent.evaluate(mock_env)
|
||||
|
||||
@patch("stable_baselines3.PPO")
|
||||
def test_evaluate_returns_metrics(
|
||||
self, mock_ppo: MagicMock, mock_env: MagicMock
|
||||
) -> None:
|
||||
"""Evaluate should return metrics dict."""
|
||||
mock_model = MagicMock()
|
||||
mock_model.predict.return_value = (np.array([0.3]), None)
|
||||
mock_ppo.return_value = mock_model
|
||||
|
||||
# Mock env.reset and env.step
|
||||
mock_env.reset.return_value = (np.random.randn(120).astype(np.float32), {})
|
||||
mock_env.step.return_value = (
|
||||
np.random.randn(120).astype(np.float32),
|
||||
0.1, # reward
|
||||
False, # terminated
|
||||
True, # truncated (end after 1 step for simplicity)
|
||||
{"return": 0.05},
|
||||
)
|
||||
|
||||
agent = RLTradingAgent(algorithm="PPO")
|
||||
agent.model = mock_model
|
||||
|
||||
metrics = agent.evaluate(mock_env, n_episodes=3)
|
||||
|
||||
assert "mean_reward" in metrics
|
||||
assert "std_reward" in metrics
|
||||
assert "mean_return" in metrics
|
||||
assert "std_return" in metrics
|
||||
assert metrics["n_episodes"] == 3
|
||||
@@ -0,0 +1,389 @@
|
||||
"""
|
||||
Tests for RL Trading Environment.
|
||||
|
||||
Covers:
|
||||
- Environment creation with various configurations
|
||||
- Observation space correctness
|
||||
- Action execution and state transitions
|
||||
- Reward calculation
|
||||
- Episode termination and truncation
|
||||
- Edge cases (empty data, extreme prices)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from typing import Tuple
|
||||
|
||||
from rdagent.components.coder.rl.env import TradingEnv, TradingState
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FIXTURES
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prices() -> np.ndarray:
|
||||
"""Generate 500-step mock price series."""
|
||||
np.random.seed(42)
|
||||
return 100.0 + np.cumsum(np.random.randn(500) * 0.5)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_indicators(mock_prices: np.ndarray) -> np.ndarray:
|
||||
"""Generate mock technical indicators (500 x 3)."""
|
||||
np.random.seed(42)
|
||||
return np.random.randn(500, 3).astype(np.float32)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def basic_env(mock_prices: np.ndarray) -> TradingEnv:
|
||||
"""Create a basic trading environment."""
|
||||
return TradingEnv(prices=mock_prices, window_size=30, max_steps=200)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def env_with_indicators(mock_prices: np.ndarray, mock_indicators: np.ndarray) -> TradingEnv:
|
||||
"""Create environment with indicators."""
|
||||
return TradingEnv(
|
||||
prices=mock_prices,
|
||||
indicators=mock_indicators,
|
||||
window_size=30,
|
||||
max_steps=200,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ENVIRONMENT CREATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestEnvCreation:
|
||||
"""Test environment initialization."""
|
||||
|
||||
def test_basic_creation(self, basic_env: TradingEnv) -> None:
|
||||
"""Environment should initialize with default values."""
|
||||
assert basic_env.initial_balance == 100000.0
|
||||
assert basic_env.transaction_cost == 0.0001
|
||||
assert basic_env.window_size == 30
|
||||
assert basic_env.current_step == 0
|
||||
assert basic_env.position == 0.0
|
||||
|
||||
def test_custom_parameters(self) -> None:
|
||||
"""Custom parameters should be respected."""
|
||||
prices = np.random.randn(200) + 100
|
||||
env = TradingEnv(
|
||||
prices=prices,
|
||||
initial_balance=50000.0,
|
||||
transaction_cost=0.0005,
|
||||
window_size=20,
|
||||
max_steps=500,
|
||||
)
|
||||
assert env.initial_balance == 50000.0
|
||||
assert env.transaction_cost == 0.0005
|
||||
assert env.window_size == 20
|
||||
assert env.max_steps == 500
|
||||
|
||||
def test_observation_space_shape_no_indicators(self, basic_env: TradingEnv) -> None:
|
||||
"""Observation dim = window_size * (1 + 0 + 3) = 30 * 4 = 120."""
|
||||
expected_dim = 30 * (1 + 0 + 3)
|
||||
assert basic_env.observation_space.shape == (expected_dim,)
|
||||
|
||||
def test_observation_space_shape_with_indicators(
|
||||
self, env_with_indicators: TradingEnv
|
||||
) -> None:
|
||||
"""Observation dim = window_size * (1 + 3 + 3) = 30 * 7 = 210."""
|
||||
expected_dim = 30 * (1 + 3 + 3)
|
||||
assert env_with_indicators.observation_space.shape == (expected_dim,)
|
||||
|
||||
def test_action_space_bounds(self, basic_env: TradingEnv) -> None:
|
||||
"""Action space should be [-1, 1]."""
|
||||
assert basic_env.action_space.low[0] == -1.0
|
||||
assert basic_env.action_space.high[0] == 1.0
|
||||
assert basic_env.action_space.shape == (1,)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# RESET
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestReset:
|
||||
"""Test environment reset."""
|
||||
|
||||
def test_reset_returns_observation_and_info(
|
||||
self, basic_env: TradingEnv
|
||||
) -> None:
|
||||
obs, info = basic_env.reset()
|
||||
assert isinstance(obs, np.ndarray)
|
||||
assert obs.dtype == np.float32
|
||||
assert isinstance(info, dict)
|
||||
|
||||
def test_reset_restores_initial_state(self, basic_env: TradingEnv) -> None:
|
||||
"""After reset, state should match initialization."""
|
||||
basic_env.reset()
|
||||
assert basic_env.current_step == 0
|
||||
assert basic_env.balance == basic_env.initial_balance
|
||||
assert basic_env.position == 0.0
|
||||
assert basic_env.entry_price == 0.0
|
||||
assert basic_env.equity_history == [basic_env.initial_balance]
|
||||
assert basic_env.trades == []
|
||||
|
||||
def test_reset_with_seed(self, mock_prices: np.ndarray) -> None:
|
||||
"""Reset with seed should be reproducible."""
|
||||
env1 = TradingEnv(prices=mock_prices, window_size=10, max_steps=50)
|
||||
env2 = TradingEnv(prices=mock_prices, window_size=10, max_steps=50)
|
||||
|
||||
obs1, _ = env1.reset(seed=123)
|
||||
obs2, _ = env2.reset(seed=123)
|
||||
np.testing.assert_array_equal(obs1, obs2)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# STEP
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestStep:
|
||||
"""Test environment step execution."""
|
||||
|
||||
def test_step_returns_correct_types(self, basic_env: TradingEnv) -> None:
|
||||
"""Step should return (obs, reward, terminated, truncated, info)."""
|
||||
basic_env.reset()
|
||||
action = np.array([0.5], dtype=np.float32)
|
||||
obs, reward, terminated, truncated, info = basic_env.step(action)
|
||||
|
||||
assert isinstance(obs, np.ndarray)
|
||||
assert isinstance(reward, float)
|
||||
assert isinstance(terminated, bool)
|
||||
assert isinstance(truncated, bool)
|
||||
assert isinstance(info, dict)
|
||||
|
||||
def test_step_updates_position(self, basic_env: TradingEnv) -> None:
|
||||
"""Position should update after step."""
|
||||
basic_env.reset()
|
||||
basic_env.step(np.array([0.5], dtype=np.float32))
|
||||
assert basic_env.position == 0.5
|
||||
|
||||
def test_step_records_trade_on_change(self, basic_env: TradingEnv) -> None:
|
||||
"""Trade should be recorded when position changes significantly."""
|
||||
basic_env.reset()
|
||||
basic_env.step(np.array([0.5], dtype=np.float32))
|
||||
assert len(basic_env.trades) == 1
|
||||
|
||||
def test_no_trade_on_small_change(self, basic_env: TradingEnv) -> None:
|
||||
"""Position changes < 0.01 should not record a trade."""
|
||||
basic_env.reset()
|
||||
basic_env.step(np.array([0.005], dtype=np.float32))
|
||||
assert len(basic_env.trades) == 0
|
||||
|
||||
def test_step_advances_time(self, basic_env: TradingEnv) -> None:
|
||||
"""current_step should increment after step."""
|
||||
basic_env.reset()
|
||||
assert basic_env.current_step == 0
|
||||
basic_env.step(np.array([0.0], dtype=np.float32))
|
||||
assert basic_env.current_step == 1
|
||||
|
||||
def test_equity_history_grows(self, basic_env: TradingEnv) -> None:
|
||||
"""Equity history should grow with each step."""
|
||||
basic_env.reset()
|
||||
initial_len = len(basic_env.equity_history)
|
||||
basic_env.step(np.array([0.5], dtype=np.float32))
|
||||
assert len(basic_env.equity_history) == initial_len + 1
|
||||
|
||||
def test_info_contains_expected_keys(self, basic_env: TradingEnv) -> None:
|
||||
"""Info dict should have standard keys."""
|
||||
basic_env.reset()
|
||||
_, _, _, _, info = basic_env.step(np.array([0.5], dtype=np.float32))
|
||||
required_keys = ["equity", "balance", "position", "trades_count", "return"]
|
||||
for key in required_keys:
|
||||
assert key in info
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# REWARD
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestReward:
|
||||
"""Test reward calculation."""
|
||||
|
||||
def test_positive_return_reward(self) -> None:
|
||||
"""Positive equity change should yield positive return component."""
|
||||
prices = np.array([100.0, 101.0, 102.0, 103.0, 104.0])
|
||||
env = TradingEnv(prices=prices, window_size=2, max_steps=10)
|
||||
env.reset()
|
||||
env.position = 1.0
|
||||
env.entry_price = 100.0
|
||||
|
||||
# Simulate positive move
|
||||
old_equity = 100000.0
|
||||
new_equity = 101000.0
|
||||
reward = env._calculate_reward(new_equity, old_equity)
|
||||
|
||||
# Return component should be positive (~0.01)
|
||||
assert reward > -0.01 # Allow small cost penalties
|
||||
|
||||
def test_negative_return_reward(self) -> None:
|
||||
"""Negative equity change should yield negative reward."""
|
||||
prices = np.array([100.0, 99.0, 98.0])
|
||||
env = TradingEnv(prices=prices, window_size=2, max_steps=10)
|
||||
env.reset()
|
||||
|
||||
old_equity = 100000.0
|
||||
new_equity = 99000.0
|
||||
reward = env._calculate_reward(new_equity, old_equity)
|
||||
assert reward < 0
|
||||
|
||||
def test_drawdown_penalty(self) -> None:
|
||||
"""Drawdown should penalize reward."""
|
||||
prices = np.array([100.0, 101.0])
|
||||
env = TradingEnv(prices=prices, window_size=2, max_steps=10)
|
||||
env.reset()
|
||||
env.equity_history = [100000.0, 110000.0, 105000.0]
|
||||
|
||||
old_equity = 105000.0
|
||||
new_equity = 104000.0
|
||||
reward_with_dd = env._calculate_reward(new_equity, old_equity)
|
||||
|
||||
# Same return without drawdown
|
||||
env.equity_history = [100000.0]
|
||||
reward_no_dd = env._calculate_reward(new_equity, old_equity)
|
||||
|
||||
assert reward_with_dd < reward_no_dd
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TERMINATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestTermination:
|
||||
"""Test episode termination conditions."""
|
||||
|
||||
def test_truncation_on_max_steps(self) -> None:
|
||||
"""Episode should truncate when max_steps reached."""
|
||||
prices = np.arange(100.0, 150.0, 0.5) # 100 steps
|
||||
env = TradingEnv(prices=prices, window_size=5, max_steps=10)
|
||||
env.reset()
|
||||
|
||||
for _ in range(10):
|
||||
obs, reward, terminated, truncated, info = env.step(np.array([0.0]))
|
||||
|
||||
assert truncated is True
|
||||
|
||||
def test_termination_on_liquidation(self) -> None:
|
||||
"""Episode should terminate on liquidation (equity < 50% initial)."""
|
||||
prices = np.array([100.0, 50.0, 10.0, 5.0, 1.0])
|
||||
env = TradingEnv(
|
||||
prices=prices,
|
||||
window_size=2,
|
||||
max_steps=10,
|
||||
initial_balance=100000.0,
|
||||
)
|
||||
env.reset()
|
||||
env.position = 1000.0 # Large long position
|
||||
env.entry_price = 100.0
|
||||
|
||||
# Price crashes -> equity drops below 50%
|
||||
obs, reward, terminated, truncated, info = env.step(np.array([1.0]))
|
||||
# May or may not terminate depending on exact equity calc
|
||||
assert isinstance(terminated, bool)
|
||||
|
||||
def test_no_termination_on_normal_step(self, basic_env: TradingEnv) -> None:
|
||||
"""Normal step should not trigger termination."""
|
||||
basic_env.reset()
|
||||
_, _, terminated, truncated, _ = basic_env.step(np.array([0.1]))
|
||||
assert terminated is False
|
||||
assert truncated is False
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# OBSERVATION
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestObservation:
|
||||
"""Test observation building."""
|
||||
|
||||
def test_observation_shape_no_indicators(self, basic_env: TradingEnv) -> None:
|
||||
"""Observation shape should match observation space."""
|
||||
obs, _ = basic_env.reset()
|
||||
assert obs.shape == basic_env.observation_space.shape
|
||||
|
||||
def test_observation_shape_with_indicators(
|
||||
self, env_with_indicators: TradingEnv
|
||||
) -> None:
|
||||
"""Observation shape should include indicator dimensions."""
|
||||
obs, _ = env_with_indicators.reset()
|
||||
assert obs.shape == env_with_indicators.observation_space.shape
|
||||
|
||||
def test_observation_values_finite(self, basic_env: TradingEnv) -> None:
|
||||
"""All observation values should be finite."""
|
||||
obs, _ = basic_env.reset()
|
||||
assert np.all(np.isfinite(obs))
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# UTILITY METHODS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestUtility:
|
||||
"""Test utility methods."""
|
||||
|
||||
def test_get_equity_curve(self, basic_env: TradingEnv) -> None:
|
||||
"""Equity curve should return array of equity values."""
|
||||
basic_env.reset()
|
||||
basic_env.step(np.array([0.5]))
|
||||
basic_env.step(np.array([0.3]))
|
||||
|
||||
curve = basic_env.get_equity_curve()
|
||||
assert isinstance(curve, np.ndarray)
|
||||
assert len(curve) == 3 # initial + 2 steps
|
||||
|
||||
def test_get_trade_log(self, basic_env: TradingEnv) -> None:
|
||||
"""Trade log should return list of trade records."""
|
||||
basic_env.reset()
|
||||
basic_env.step(np.array([0.5]))
|
||||
basic_env.step(np.array([-0.3]))
|
||||
|
||||
log = basic_env.get_trade_log()
|
||||
assert len(log) == 2
|
||||
assert "step" in log[0]
|
||||
assert "action" in log[0]
|
||||
assert "cost" in log[0]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TRADING STATE DATACLASS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestTradingState:
|
||||
"""Test TradingState dataclass."""
|
||||
|
||||
def test_default_values(self) -> None:
|
||||
"""Default values should match initialization params."""
|
||||
state = TradingState()
|
||||
assert state.position == 0.0
|
||||
assert state.cash == 100000.0
|
||||
assert state.equity == 100000.0
|
||||
assert state.entry_price == 0.0
|
||||
assert state.step == 0
|
||||
assert state.holdings_history == []
|
||||
|
||||
def test_custom_values(self) -> None:
|
||||
"""Custom values should be stored correctly."""
|
||||
state = TradingState(
|
||||
position=0.5,
|
||||
cash=50000.0,
|
||||
equity=75000.0,
|
||||
entry_price=100.0,
|
||||
step=10,
|
||||
holdings_history=[0.1, 0.2, 0.3],
|
||||
)
|
||||
assert state.position == 0.5
|
||||
assert state.cash == 50000.0
|
||||
assert len(state.holdings_history) == 3
|
||||
Reference in New Issue
Block a user