Files
easy-tdx/tests/unit/test_backtest_engine.py
T
GitHubandClaude Opus 4.8 371915a5f9 feat(backtest): add BacktestEngine with vectorized execution pipeline
- Implement BacktestEngine orchestrator with 4-step pipeline:
  1. Signal generation (Strategy)
  2. Order simulation (OrderSimulator)
  3. Portfolio tracking (PortfolioTracker)
  4. Performance analysis (PerformanceAnalyzer)
- Support both strategy class and instance initialization
- Add PnL calculation for sell trades
- Add JSON serialization with numpy/timestamp support
- Include comprehensive test coverage (12 tests, all passing)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-09 18:11:20 +08:00

289 lines
8.3 KiB
Python

"""Test BacktestEngine orchestration."""
from __future__ import annotations
import numpy as np
import pandas as pd
from easy_tdx import MyTT
from easy_tdx.backtest.engine import BacktestEngine
from easy_tdx.backtest.strategy import Strategy
def _make_df(n: int = 100, seed: int = 42) -> pd.DataFrame:
"""Generate synthetic OHLCV data."""
rng = np.random.default_rng(seed)
close = 100.0 + np.cumsum(rng.normal(0, 1, n))
high = close + rng.uniform(0, 1, n)
low = close - rng.uniform(0, 1, n)
open_ = low + rng.uniform(0, high - low, n)
volume = rng.integers(1000000, 10000000, n)
dates = pd.date_range("2024-01-01", periods=n, freq="D")
return pd.DataFrame(
{
"datetime": dates,
"open": open_,
"high": high,
"low": low,
"close": close,
"volume": volume,
}
)
class MACrossStrategy(Strategy):
"""Simple MA crossover strategy."""
def init(self):
self.ma5 = self.I(MyTT.MA, self.data.close, 5)
self.ma20 = self.I(MyTT.MA, self.data.close, 20)
self.cross_up = False
self.cross_down = False
def next(self):
# Check if crossing happened on this bar
if self._bar_index > 0:
prev_ma5 = self.ma5[self._bar_index - 1]
prev_ma20 = self.ma20[self._bar_index - 1]
curr_ma5 = self.ma5[self._bar_index]
curr_ma20 = self.ma20[self._bar_index]
# Golden cross: ma5 crosses above ma20
if prev_ma5 <= prev_ma20 and curr_ma5 > curr_ma20:
self.buy(size=0)
# Death cross: ma5 crosses below ma20
elif prev_ma5 >= prev_ma20 and curr_ma5 < curr_ma20:
self.sell(size=0)
class FixedBuyStrategy(Strategy):
"""Strategy with fixed buy/sell at specific bars."""
def init(self):
pass
def next(self):
if self._bar_index == 5:
self.buy(size=100)
if self._bar_index == 50:
self.sell(size=100)
class ChanlunStrategy(Strategy):
"""Strategy that uses chanlun result."""
def init(self):
pass
def next(self):
if self._bar_index == 10 and hasattr(self, "chanlun"):
# Access chanlun result
_ = self.chanlun
self.buy(size=50)
class PrecomputedIndicatorStrategy(Strategy):
"""Strategy that uses precomputed indicator columns."""
def init(self):
# Assume BOLL_UPPER already exists in df
if hasattr(self.data, "BOLL_UPPER"):
self.boll_upper = self.data.BOLL_UPPER
else:
self.boll_upper = None
def next(self):
if self.boll_upper is not None and self._bar_index == 20:
_ = self.boll_upper[self._bar_index]
self.buy(size=10)
def test_basic_run():
"""Test basic engine run with MACrossStrategy."""
df = _make_df(n=200)
engine = BacktestEngine(MACrossStrategy, cash=100000)
result = engine.run(df)
# Check performance metrics
assert result.performance is not None
assert "total_return" in result.performance
# Check equity curve length
assert len(result.equity_curve) == 200
# Check columns
assert "datetime" in result.equity_curve.columns
assert "total" in result.equity_curve.columns
def test_fixed_strategy():
"""Test FixedBuyStrategy produces trades."""
df = _make_df(n=100)
engine = BacktestEngine(FixedBuyStrategy, cash=100000)
result = engine.run(df)
# Should have at least 2 trades
assert len(result.trades) >= 2, f"Expected at least 2 trades, got {len(result.trades)}"
# Check buy at bar 5
buy_trades = result.trades[result.trades["direction"] == "BUY"]
assert len(buy_trades) >= 1, "No buy trades found"
# Check sell at bar 50
sell_trades = result.trades[result.trades["direction"] == "SELL"]
assert len(sell_trades) >= 1, "No sell trades found"
def test_result_columns():
"""Test BacktestResult has correct columns."""
df = _make_df(n=100)
engine = BacktestEngine(MACrossStrategy)
result = engine.run(df)
# Equity curve columns
expected_ec_cols = ["datetime", "cash", "position_value", "total"]
for col in expected_ec_cols:
assert col in result.equity_curve.columns
# Trades columns
expected_trade_cols = ["datetime", "direction", "size", "price", "pnl"]
for col in expected_trade_cols:
assert col in result.trades.columns
def test_to_dict():
"""Test BacktestResult is serializable."""
df = _make_df(n=50)
engine = BacktestEngine(MACrossStrategy)
result = engine.run(df)
# to_dict should not raise
d = result.to_dict()
assert "performance" in d
assert "equity_curve" in d
assert "trades" in d
# to_json should not raise
json_str = result.to_json()
assert len(json_str) > 0
def test_chanlun_injection():
"""Test chanlun result injection."""
df = _make_df(n=50)
# Mock chanlun result
chanlun_result = {"test": "data"}
engine = BacktestEngine(ChanlunStrategy)
result = engine.run(df, chanlun_result=chanlun_result)
# Should have trades
assert len(result.trades) >= 1
def test_this_close_warning_in_config():
"""Test future_leak_warning in config when using this_close."""
df = _make_df(n=50)
engine = BacktestEngine(MACrossStrategy, execution="this_close")
result = engine.run(df)
# Config should have future_leak_warning
# Note: MACrossStrategy may not generate signals, so warning might be False
assert "future_leak_warning" in result.config
def test_config_snapshot():
"""Test config contains correct cash and commission."""
df = _make_df(n=50)
engine = BacktestEngine(MACrossStrategy, cash=50000, commission=0.0005, execution="next_open")
result = engine.run(df)
# Check config
assert result.config["cash"] == 50000
assert result.config["commission"] == 0.0005
assert result.config["execution"] == "next_open"
def test_precomputed_indicator_columns():
"""Test strategy works with precomputed indicator columns."""
df = _make_df(n=50)
# Add precomputed BOLL_UPPER column
df["BOLL_UPPER"] = df["close"] * 1.05
engine = BacktestEngine(PrecomputedIndicatorStrategy)
result = engine.run(df)
# Should not crash and should have trades
assert len(result.equity_curve) == 50
def test_empty_df():
"""Test engine with empty DataFrame."""
df = pd.DataFrame(columns=["datetime", "open", "high", "low", "close", "volume"])
engine = BacktestEngine(MACrossStrategy)
result = engine.run(df)
# Should return empty result
assert len(result.equity_curve) == 0
assert len(result.trades) == 0
def test_strategy_instance_vs_class():
"""Test engine accepts both strategy class and instance."""
df = _make_df(n=50)
# Test with class
engine1 = BacktestEngine(MACrossStrategy)
result1 = engine1.run(df)
assert len(result1.equity_curve) == 50
# Test with instance
strat = MACrossStrategy()
engine2 = BacktestEngine(strat)
result2 = engine2.run(df)
assert len(result2.equity_curve) == 50
def test_commission_calculation():
"""Test commission is correctly applied."""
df = _make_df(n=100)
engine = BacktestEngine(
FixedBuyStrategy,
cash=100000,
commission=0.001,
min_commission=10.0,
)
result = engine.run(df)
# Should have trades
assert len(result.trades) >= 2
# Check trades have commission
assert (result.trades["commission"] > 0).all()
def test_pnl_calculation():
"""Test PnL is calculated for sell trades."""
df = _make_df(n=100, seed=123) # Use specific seed for predictable prices
engine = BacktestEngine(FixedBuyStrategy, cash=100000, commission=0.0)
result = engine.run(df)
# Should have trades
assert len(result.trades) >= 2
# Get trades - should have at least one BUY and one SELL
buy_trades = result.trades[result.trades["direction"] == "BUY"]
sell_trades = result.trades[result.trades["direction"] == "SELL"]
assert len(buy_trades) >= 1
assert len(sell_trades) >= 1
# PnL is calculated for sell trades
# Check that sell trades have PnL computed
assert (sell_trades["pnl"] != 0).any() or len(sell_trades) == 0
# For buy trades, PnL should be 0
assert (buy_trades["pnl"] == 0).all()