import unittest import json import tempfile from pathlib import Path from unittest import mock import numpy as np import pandas as pd import data_manager import trader from execution import simulate_book, rebalance, portfolio_value from execution import compute_trades, market_permissions from signals import session_targets from strategies.momentum import MomentumStrategy from strategies.momentum_quality import rolling_max_drawdown from strategies.composite_alpha import CompositeAlphaStrategy from strategies.robust_momentum import RobustMomentumStrategy class ExecutionContracts(unittest.TestCase): def test_a_halt_cannot_be_sold_at_the_carried_mark(self): dates = pd.bdate_range("2026-01-05", periods=5) raw = pd.DataFrame({"A": [100, np.nan, np.nan, 10, 10]}, index=dates) targets = pd.DataFrame({"A": [1, 0, 0, 0, 0]}, index=dates) result = simulate_book(targets, raw, 10000, transaction_cost=0, min_trade_value=0) self.assertAlmostEqual(result["equity"].iloc[-1], 1000) self.assertEqual([t["date"] for t in result["trade_log"]], [str(dates[0].date()), str(dates[3].date())]) def test_fractional_and_integer_orders_are_funded_after_all_fees(self): prices = {"A": 73., "B": 29.} for integer in [False, True]: state = {"holdings": {}, "cash": 1000., "trade_log": []} rebalance(state, {"A": .5, "B": .5}, prices, set(prices), fixed_fee=2, fee_base=1.88, fee_per_share=.009, integer_shares=integer, min_trade_value=0) self.assertGreaterEqual(state["cash"], 0) equity = portfolio_value(state["holdings"], prices, state["cash"]) self.assertAlmostEqual(equity, 1000 - sum(t["commission"] for t in state["trade_log"]), places=2) def test_batch_replay_matches_session_ledger_with_a_halt_and_drift(self): dates = pd.bdate_range("2026-01-05", periods=5) prices = pd.DataFrame({"A": [100, 110, np.nan, 70, 90], "B": [20, 21, 22, 22, 23]}, index=dates) targets = pd.DataFrame([[.5, .5], [.5, .5], [0, 1], [0, 1], [.5, .5]], index=dates, columns=prices.columns) options = dict(transaction_cost=.001, fixed_fee=2., integer_shares=True, resize_band=.1) batch = simulate_book(targets, prices, 10000, **options) state = trader.init_state("us", "test", 10000) fees = {"tx_cost": .001, "fixed_fee": 2.} for day in dates: trader._run_session(state, "us", day, "B", prices.ffill(), prices, fees, True, targets.loc[day]) self.assertAlmostEqual(state["daily_equity"][str(day.date())], batch["equity"].loc[day], places=2) self.assertEqual(state["trade_log"], batch["trade_log"]) def test_a_delisted_stock_is_settled_at_its_last_print_in_both_ledgers(self): dates = pd.bdate_range("2026-01-05", periods=9) # A is acquired: last print on day 2, out of the index from day 4. # H halts on the same day but stays a member, so it stays frozen. prices = pd.DataFrame({"A": [100, 101, 102] + [np.nan] * 6, "H": [50, 50, 51] + [np.nan] * 6, "B": [20.] * 9}, index=dates) member = pd.DataFrame(True, index=dates, columns=prices.columns) member.loc[dates[4]:, "A"] = False targets = pd.DataFrame([[.4, .3, .3]] * 9, index=dates, columns=prices.columns) options = dict(transaction_cost=.001, fixed_fee=2., integer_shares=True, resize_band=.1) batch = simulate_book(targets, prices, 10000, market="us", member=member, **options) settled = [(t["ticker"], t["date"], t["price"]) for t in batch["trade_log"] if t.get("settled")] self.assertEqual(settled, [("A", str(dates[7].date()), 102.0)]) self.assertIn("H", batch["holdings"]) state = trader.init_state("us", "test", 10000) fees = {"tx_cost": .001, "fixed_fee": 2.} for day in dates: trader._run_session(state, "us", day, "B", prices.ffill(), prices, fees, True, targets.loc[day], member_row=member.loc[day]) self.assertAlmostEqual(state["daily_equity"][str(day.date())], batch["equity"].loc[day], places=2) self.assertEqual(state["trade_log"], batch["trade_log"]) china = simulate_book(targets, prices, 10000, market="cn", member=member, **options) self.assertFalse(any(t.get("settled") for t in china["trade_log"])) def test_contributions_are_executed_and_fees_affect_unit_returns(self): dates = pd.bdate_range("2026-01-05", periods=3) prices = pd.DataFrame({"A": [100., 100., 100.]}, index=dates) weights = pd.DataFrame(1., index=dates, columns=["A"]) base = simulate_book(weights, prices, 1000, transaction_cost=0, fixed_fee=2) injected = simulate_book(weights, prices, 1000, transaction_cost=0, fixed_fee=2, contributions=pd.Series([0, 1000, 0], index=dates)) self.assertEqual(len(base["trade_log"]), 1) self.assertEqual(len(injected["trade_log"]), 2) self.assertAlmostEqual(injected["equity"].iloc[-1], 1996) self.assertLess(injected["unit_equity"].iloc[-1], base["unit_equity"].iloc[-1]) def test_missing_open_never_sees_same_day_close(self): close = pd.DataFrame({"A": [100., 200.]}) raw_open = pd.DataFrame({"A": [np.nan, np.nan]}) self.assertTrue(data_manager._fill_open(close, raw_open).isna().all().all()) def test_china_limit_permissions_tax_and_star_minimum(self): raw = pd.Series({"600000.SS": 11., "000001.SZ": 9., "688001.SS": 12.}) buy, sell, tax = market_permissions("cn", pd.Timestamp("2026-09-30"), raw, raw * 0 + 10) self.assertNotIn("600000.SS", buy) self.assertNotIn("000001.SZ", sell) self.assertNotIn("688001.SS", buy) self.assertEqual(tax, .0005) self.assertEqual(market_permissions("cn", pd.Timestamp("2023-08-25"), raw, raw * 0 + 10)[2], .001) orders = compute_trades({}, 10000, {"688001.SS": .7}, {"688001.SS": 23.}, integer_shares=True, lot_size=100) self.assertEqual(orders[0]["shares_delta"], 304) class SignalContracts(unittest.TestCase): def test_live_refresh_updates_the_membership_snapshot_before_selection(self): with tempfile.TemporaryDirectory() as tmp: cache = Path(tmp) / "sp500_history.json" cache.write_text(json.dumps({"date": "2000-01-01", "intervals": {"OLD": [["2020-01-01", None]]}})) history = {"OLD": [["2020-01-01", "2026-09-30"]], "NEW": [["2026-09-30", None]]} with mock.patch.object(trader.universe_history, "CACHE_DIR", tmp), \ mock.patch.object(trader.universe_history, "build_sp500_history", return_value=history), \ mock.patch.object(trader, "strategy_universe", return_value=(["NEW"], "SPY")), \ mock.patch.object(data_manager, "update"): trader.refresh_panels("us", ["momentum"], "test") member = trader.index_members("us", pd.DatetimeIndex(["2026-10-01"])) self.assertTrue(member.iloc[0]["NEW"]) self.assertFalse(member.iloc[0]["OLD"]) def test_close_signal_is_unaffected_by_its_execution_days_close(self): dates = pd.bdate_range("2024-01-01", periods=30) c = pd.DataFrame({"A": np.arange(100., 130.), "B": np.arange(130., 100., -1)}, index=dates) s = MomentumStrategy(lookback=5, skip=1, top_n=1) base = session_targets(s, c) shocked = c.copy(); shocked.iloc[20:] = shocked.iloc[20:] * [0.1, 10] actual = session_targets(s, shocked) pd.testing.assert_frame_equal(base.iloc[:21], actual.iloc[:21]) def test_trader_and_backtest_use_identical_membership_histories(self): rng = np.random.default_rng(42) dates = pd.bdate_range("2024-01-01", periods=370) close = pd.DataFrame(100 * np.exp(rng.normal(.0005, .02, (370, 25)).cumsum(0)), index=dates, columns=[f"S{i}" for i in range(25)]) opens = close.shift(1).fillna(100) member = pd.DataFrame(True, index=dates, columns=close.columns) member.iloc[:320, 0] = False member.iloc[330:, 1] = False bars = {"close": close, "open": opens} strategy = RobustMomentumStrategy("overnight_recovery", top_n=5) expected = session_targets(strategy, close, opens, bars, member) with mock.patch.object(trader, "index_members", lambda m, d: member.reindex(d)): actual = trader.session_signals("us", "overnight_recovery_top10", list(close), dates[-20:], close, opens, bars=bars, strategy=strategy) pd.testing.assert_frame_equal(actual, expected.reindex(index=actual.index, columns=actual.columns)) self.assertTrue((actual["S1"] == 0).all()) def test_quality_prefers_a_smaller_drawdown_and_needs_only_one_window(self): p = pd.DataFrame({"small": [100, 110, 99, 108], "large": [100, 110, 55, 108]}) dd = rolling_max_drawdown(p, 4) self.assertGreater(dd.iloc[-1]["small"], dd.iloc[-1]["large"]) self.assertAlmostEqual(dd.iloc[-1]["small"], -.1) def test_missing_regime_benchmark_is_an_error(self): with self.assertRaisesRegex(ValueError, "benchmark"): CompositeAlphaStrategy().generate_signals(pd.DataFrame({"A": [1, 2]})) if __name__ == "__main__": unittest.main()