"""Pure-function tests for the risk guardrails — no network, deterministic.""" from __future__ import annotations from datetime import datetime, timezone from agents.schemas import AgentOutput, Pick from config.risk import RiskConfig from execution.guardrails import plan_orders EQUITY = 100_000.0 CAPS = {"nova": 0.10, "warren_lite": 0.08, "atlas": 0.09} def _out(persona_id: str, picks: list[Pick]) -> AgentOutput: return AgentOutput( agent_name=persona_id, persona_id=persona_id, started_at=datetime.now(timezone.utc), picks=picks, ) def _pick(ticker, conviction=0.7, weight=0.08, direction="long") -> Pick: return Pick( ticker=ticker, conviction=conviction, thesis="t", requested_weight_pct=weight, direction=direction, ) def _plan(agent_outputs, prices, *, risk=None, held=None, buying_power=EQUITY): return plan_orders( agent_outputs, equity=EQUITY, buying_power=buying_power, prices=prices, held_symbols=held or set(), risk=risk or RiskConfig(), persona_caps=CAPS, ) def test_shorts_dropped_when_disabled(): orders, notes = _plan( [_out("nova", [_pick("NVDA"), _pick("TSLA", direction="short")])], {"NVDA": 200.0, "TSLA": 250.0}, risk=RiskConfig(allow_shorts=False), ) tickers = {o.ticker for o in orders} assert "NVDA" in tickers and "TSLA" not in tickers assert any("shorts disabled" in n for n in notes) def test_shorts_kept_when_enabled(): orders, _ = _plan( [_out("nova", [_pick("NVDA"), _pick("TSLA", direction="short")])], {"NVDA": 200.0, "TSLA": 250.0}, risk=RiskConfig(allow_shorts=True), ) by_ticker = {o.ticker: o for o in orders} assert "TSLA" in by_ticker and by_ticker["TSLA"].direction == "short" assert by_ticker["NVDA"].direction == "long" def test_dedup_highest_conviction_wins(): orders, _ = _plan( [ _out("nova", [_pick("MSFT", conviction=0.6)]), _out("atlas", [_pick("MSFT", conviction=0.9)]), ], {"MSFT": 400.0}, ) msft = [o for o in orders if o.ticker == "MSFT"] assert len(msft) == 1 assert msft[0].persona_id == "atlas" and msft[0].conviction == 0.9 def test_per_agent_trade_cap(): picks = [_pick(t, conviction=c) for t, c in [("A", 0.9), ("B", 0.8), ("C", 0.7), ("D", 0.6)]] orders, notes = _plan( [_out("nova", picks)], {"A": 10.0, "B": 10.0, "C": 10.0, "D": 10.0}, risk=RiskConfig(max_trades_per_agent=2), ) assert len(orders) == 2 assert {o.ticker for o in orders} == {"A", "B"} # top conviction kept def test_min_price_filter(): orders, notes = _plan([_out("nova", [_pick("PENNY")])], {"PENNY": 1.5}) assert orders == [] assert any("min" in n for n in notes) def test_skip_already_held(): orders, notes = _plan( [_out("nova", [_pick("NVDA"), _pick("AMD")])], {"NVDA": 200.0, "AMD": 150.0}, held={"NVDA"}, ) assert {o.ticker for o in orders} == {"AMD"} assert any("already held" in n for n in notes) def test_sizing_and_weight_clamp(): # requested 50% clamps to nova cap 10% -> qty = floor(0.10 * 100000 / 200) = 50 orders, _ = _plan([_out("nova", [_pick("NVDA", weight=0.50)])], {"NVDA": 200.0}) assert len(orders) == 1 o = orders[0] assert abs(o.weight_pct - 0.10) < 1e-9 assert o.qty == 50 def test_day_cap_truncates_by_conviction(): picks = [_pick(f"T{i}", conviction=0.5 + i / 100) for i in range(6)] orders, notes = _plan( [_out("nova", picks[:3]), _out("atlas", picks[3:])], {f"T{i}": 100.0 for i in range(6)}, risk=RiskConfig(max_trades_per_day=3, max_trades_per_agent=5), ) assert len(orders) == 3 assert any("day cap" in n for n in notes) def test_option_sizing_by_premium(): opt = Pick( ticker="AAPL", conviction=0.8, thesis="bullish", requested_weight_pct=0.10, instrument="option", option_right="call", contract_symbol="AAPL260724C00295000", strike=295.0, expiry="2026-07-24", ) orders, _ = _plan([_out("nova", [opt])], {"AAPL260724C00295000": 7.39}) assert len(orders) == 1 o = orders[0] assert o.instrument == "option" and o.multiplier == 100 assert o.contract_symbol == "AAPL260724C00295000" # premium cap: weight = min(requested, persona cap, max_option_weight_pct=0.10) # qty = floor(0.10 * 100000 / (7.39 * 100)) = floor(10000 / 739) = 13 assert abs(o.weight_pct - 0.10) < 1e-9 assert o.qty == 13 def test_buying_power_trim(): # 2 orders needing 50k each but only 60k buying power -> keep the higher conviction one orders, notes = _plan( [_out("nova", [_pick("HI", conviction=0.9, weight=0.10), _pick("LO", conviction=0.5, weight=0.10)])], {"HI": 200.0, "LO": 200.0}, buying_power=12_000.0, # only one 50-share*200=10k order fits ) assert len(orders) == 1 assert orders[0].ticker == "HI" assert any("buying power" in n for n in notes) def _opt_pick(ticker="AAPL", conviction=0.8, weight=0.05, right="call"): sym = f"{ticker}260724{'C' if right == 'call' else 'P'}00295000" return Pick( ticker=ticker, conviction=conviction, thesis="t", requested_weight_pct=weight, instrument="option", option_right=right, contract_symbol=sym, strike=295.0, expiry="2026-07-24", ) def test_option_conviction_gate(): lo = _opt_pick("AAPL", conviction=0.4) hi = _opt_pick("MSFT", conviction=0.8) orders, notes = _plan( [_out("nova", [lo, hi])], {lo.contract_symbol: 3.0, hi.contract_symbol: 3.0}, ) assert {o.ticker for o in orders} == {"MSFT"} assert any("conviction" in n and "floor" in n for n in notes) def test_option_reserved_in_per_agent_cap(): # option is the LOWEST conviction of 5 picks with cap 4 — it must still survive eq = [_pick(t, conviction=c) for t, c in [("A", 0.9), ("B", 0.85), ("C", 0.8), ("D", 0.75)]] opt = _opt_pick("AAPL", conviction=0.7) prices = {p.ticker: 100.0 for p in eq} | {opt.contract_symbol: 3.0} orders, notes = _plan([_out("nova", [*eq, opt])], prices) assert len(orders) == 4 assert any(o.instrument == "option" for o in orders) assert any("reserved option slot" in n for n in notes) def test_option_reserved_in_day_cap(): # day cap 2, option ranks 3rd by conviction -> swapped back in eq = [_pick(t, conviction=c) for t, c in [("A", 0.9), ("B", 0.85)]] opt = _opt_pick("AAPL", conviction=0.7) prices = {p.ticker: 100.0 for p in eq} | {opt.contract_symbol: 3.0} orders, notes = _plan( [_out("nova", [*eq, opt])], prices, risk=RiskConfig(max_trades_per_day=2, max_trades_per_agent=5), ) assert len(orders) == 2 assert any(o.instrument == "option" for o in orders) assert any("reserved option" in n for n in notes) def test_option_survives_buying_power_trim(): # option has lowest conviction but equities get dropped first eq = [_pick("BIG", conviction=0.9, weight=0.10)] opt = _opt_pick("AAPL", conviction=0.7, weight=0.05) prices = {"BIG": 200.0, opt.contract_symbol: 5.0} # equity order = 50 * 200 = 10k; option = qty 10 * 5 * 100 = 5k; BP only fits the option orders, notes = _plan([_out("nova", [*eq, opt])], prices, buying_power=6_000.0) assert len(orders) == 1 assert orders[0].instrument == "option"