"""Unit tests for ORBTradingEngine._compute_sizing_capital.""" from __future__ import annotations from types import SimpleNamespace from unittest.mock import MagicMock import pytest from apps.orb_trader.engine import ORBTradingEngine def _make_engine( *, initial_equity: float = 10_000.0, daily_budget_reset: bool = True, drawdown_governor_threshold: float | None = None, drawdown_governor_min_scale: float = 0.30, streak_sizing_win_bonus: float | None = None, streak_sizing_loss_penalty: float | None = None, streak_sizing_max: float = 2.5, streak_sizing_min: float = 0.5, peak_equity: float | None = None, trades: list[dict] | None = None, ) -> ORBTradingEngine: params = SimpleNamespace( daily_budget_reset=daily_budget_reset, drawdown_governor_threshold=drawdown_governor_threshold, drawdown_governor_min_scale=drawdown_governor_min_scale, streak_sizing_win_bonus=streak_sizing_win_bonus, streak_sizing_loss_penalty=streak_sizing_loss_penalty, streak_sizing_max=streak_sizing_max, streak_sizing_min=streak_sizing_min, ) session = SimpleNamespace( session_id="test-session", initial_equity=initial_equity, ) state = MagicMock() state.get_peak_equity.return_value = peak_equity if peak_equity is not None else initial_equity state.list_trades.return_value = trades or [] engine = object.__new__(ORBTradingEngine) engine._params = params engine._session = session engine._state = state return engine # ── daily_budget_reset ──────────────────────────────────────────────────────── class TestDailyBudgetReset: def test_reset_true_uses_initial_equity(self): eng = _make_engine(initial_equity=10_000, daily_budget_reset=True) assert eng._compute_sizing_capital(15_000) == pytest.approx(10_000) def test_reset_false_uses_current_equity(self): eng = _make_engine(initial_equity=10_000, daily_budget_reset=False) assert eng._compute_sizing_capital(15_000) == pytest.approx(15_000) # ── Drawdown governor ───────────────────────────────────────────────────────── class TestDrawdownGovernor: def test_no_drawdown_no_scaling(self): eng = _make_engine( drawdown_governor_threshold=0.025, drawdown_governor_min_scale=0.30, peak_equity=10_000, ) result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000) def test_drawdown_below_threshold_no_scaling(self): # 2% DD, threshold 2.5% → no scaling eng = _make_engine( drawdown_governor_threshold=0.025, peak_equity=10_000, ) result = eng._compute_sizing_capital(9_800) assert result == pytest.approx(10_000) def test_drawdown_at_full_governor(self): # DD = 5% = 2 * threshold(2.5%) → excess = 1x threshold → scale = min_scale eng = _make_engine( drawdown_governor_threshold=0.025, drawdown_governor_min_scale=0.30, peak_equity=10_000, ) result = eng._compute_sizing_capital(9_500) assert result == pytest.approx(10_000 * 0.30) def test_drawdown_partial_governor(self): # DD = 3.75% → excess = 1.25% = 0.5 * threshold(2.5%) # scale = max(0.30, 1.0 - 0.70 * 0.5) = max(0.30, 0.65) = 0.65 eng = _make_engine( drawdown_governor_threshold=0.025, drawdown_governor_min_scale=0.30, peak_equity=10_000, ) result = eng._compute_sizing_capital(9_625) expected_scale = max(0.30, 1.0 - 0.70 * 0.5) assert result == pytest.approx(10_000 * expected_scale, rel=1e-4) # ── Streak sizing ───────────────────────────────────────────────────────────── class TestStreakSizing: def _win(self, pnl=100.0) -> dict: return {"pnl": pnl} def _loss(self, pnl=-100.0) -> dict: return {"pnl": pnl} def test_no_trades_no_multiplier(self): eng = _make_engine(streak_sizing_win_bonus=0.70, trades=[]) assert eng._compute_sizing_capital(10_000) == pytest.approx(10_000) def test_single_win_streak_1(self): # streak_len=1 win → mult = 1 + 1*0.70 = 1.70 eng = _make_engine( streak_sizing_win_bonus=0.70, trades=[self._win()], # newest first ) result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000 * 1.70) def test_two_wins_streak_2(self): # streak_len=2 → mult = 1 + 2*0.70 = 2.40 eng = _make_engine( streak_sizing_win_bonus=0.70, trades=[self._win(), self._win()], ) result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000 * 2.40) def test_win_streak_capped_at_max(self): # streak_len=5 → mult = 1 + 5*0.70 = 4.50 → capped at streak_max=2.5 eng = _make_engine( streak_sizing_win_bonus=0.70, streak_sizing_max=2.5, trades=[self._win()] * 5, ) result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000 * 2.5) def test_loss_streak_reduces_sizing(self): # streak_len=2 loss, loss_penalty=0.20 → mult = 1 - 2*0.20 = 0.60 eng = _make_engine( streak_sizing_loss_penalty=0.20, streak_sizing_min=0.5, trades=[self._loss(), self._loss()], ) result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000 * 0.60) def test_loss_streak_floored_at_min(self): eng = _make_engine( streak_sizing_loss_penalty=0.20, streak_sizing_min=0.5, trades=[self._loss()] * 10, ) result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000 * 0.5) def test_streak_direction_newest_first(self): # trades list DESC (newest first): [win, loss, loss] # The most recent is a win → streak_len=1 → mult=1.70 eng = _make_engine( streak_sizing_win_bonus=0.70, trades=[self._win(), self._loss(), self._loss()], ) result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000 * 1.70) def test_streak_direction_oldest_not_used(self): # trades list DESC: [loss, win, win] # Most recent = loss, streak_len=1 → no win_bonus applies (only loss_penalty) eng = _make_engine( streak_sizing_win_bonus=0.70, trades=[self._loss(), self._win(), self._win()], ) # No loss_penalty configured, so streak_mult = 1.0 result = eng._compute_sizing_capital(10_000) assert result == pytest.approx(10_000 * 1.0)