"""Contracts added with the robust-momentum strategies: bar strategies are causal and trade on the live trader's timing, whole-lot orders track targets, panel backfills retry, and point-in-time CSI 300 membership.""" import unittest from unittest import mock import numpy as np import pandas as pd import data_manager import trader import universe_history as uh from main import backtest, simulate_live_book, strategy_positions from strategies.robust_momentum import SIGNALS, RobustMomentumStrategy, select_with_buffer def _bars(n_days: int = 420, n_assets: int = 30, seed: int = 0) -> dict[str, pd.DataFrame]: rng = np.random.default_rng(seed) idx = pd.bdate_range("2024-01-01", periods=n_days) cols = [f"S{i:02d}" for i in range(n_assets)] r = rng.normal(0.0004, 0.02, (n_days, n_assets)) close = pd.DataFrame(100 * np.exp(np.cumsum(r, axis=0)), index=idx, columns=cols) gap = rng.normal(0, 0.005, (n_days, n_assets)) opens = close.shift(1).bfill() * np.exp(gap) high = np.maximum(opens, close) * (1 + rng.uniform(0, 0.02, (n_days, n_assets))) low = np.minimum(opens, close) * (1 - rng.uniform(0, 0.02, (n_days, n_assets))) volume = pd.DataFrame(rng.uniform(1e5, 1e6, (n_days, n_assets)), index=idx, columns=cols) return {"close": close, "open": opens, "high": high, "low": low, "volume": volume} class RobustMomentumTests(unittest.TestCase): def test_weights_on_a_day_ignore_that_day_and_later_bars(self): bars = _bars() cut = bars["close"].index[380] for signal in SIGNALS: strat = RobustMomentumStrategy(signal, top_n=5, rebal_freq=5, buffer=1.0) base = strat.generate_signals(bars) shocked = {f: b.copy() for f, b in bars.items()} for f in shocked: shocked[f].loc[cut:] *= 3.0 again = strat.generate_signals(shocked) pd.testing.assert_frame_equal(base.loc[:cut], again.loc[:cut], obj=signal) self.assertGreater(base.loc[cut].sum(), 0.99, signal) def test_held_names_stay_while_inside_the_buffer(self): idx = pd.bdate_range("2025-01-01", periods=3) score = pd.DataFrame([[5, 4, 3, 2, 1], [1, 2, 5, 4, 3], [1, 2, 5, 4, 3]], index=idx, columns=list("ABCDE"), dtype=float) w = select_with_buffer(score, top_n=2, rebal_freq=1, buffer=1.0, warmup=0) self.assertEqual(set(w.columns[w.iloc[0] > 0]), {"A", "B"}) # Day 2: A ranks 5th (outside 2 * (1 + 1) = 4) and leaves; B ranks 4th and stays. self.assertEqual(set(w.columns[w.iloc[1] > 0]), {"B", "C"}) def test_bar_strategy_trades_at_the_close_after_its_signal_day(self): bars = _bars() strat = RobustMomentumStrategy("lowbeta_sharpe_mom", top_n=5, rebal_freq=5) weights = strat.generate_signals(bars) held = strategy_positions(strat, bars["close"], open_data=bars["open"], bars=bars) pd.testing.assert_frame_equal(held, weights.shift(1).fillna(0.0)) equity = backtest(strat, bars["close"], initial_capital=1000, open_data=bars["open"], bars=bars) self.assertEqual(len(equity), len(bars["close"])) def test_trader_session_signals_use_bars_through_the_previous_day(self): bars = _bars() days = bars["close"].index[-3:] strat = RobustMomentumStrategy("overnight_recovery", top_n=5, rebal_freq=1) full = strat.generate_signals(bars).loc[days] got = trader.bar_session_weights(strat, bars, days) pd.testing.assert_frame_equal(got, full) # Morning run: today's bar does not exist yet. today = days[-1] + pd.offsets.BDay() morning = trader.bar_session_weights(strat, bars, pd.DatetimeIndex([today])) self.assertAlmostEqual(morning.loc[today].sum(), 1.0) class LiveBookSimulationTests(unittest.TestCase): def test_holdings_drift_and_only_resizes_beyond_the_band_trade(self): idx = pd.bdate_range("2026-01-05", periods=5) data = pd.DataFrame({"A": [100, 100, 125, 125, 125], "B": [100] * 5, "C": [100] * 5}, index=idx, dtype=float) third = [0, 1 / 3, 1 / 3, 1 / 3] held = pd.DataFrame({"A": third + [0.0], "B": third + [0.5], "C": third + [0.5]}, index=idx) equity, traded = simulate_live_book(held, data, 1000.0, transaction_cost=0.0, resize_band=0.10) orders = (traded.abs() > 0).sum(axis=1).tolist() # Buy all three; A +25% drifts to 38.5% (15% over target): trim A only, # B and C at 30.8% stay inside the band; then exit A, top up B and C. self.assertEqual(orders, [0, 3, 0, 1, 3]) self.assertLess(traded.loc[idx[3], "A"], 0) self.assertAlmostEqual(equity.iloc[2], 1000 * (1 + 0.25 / 3), places=6) def test_orders_below_the_floor_are_skipped_but_exits_are_not(self): idx = pd.bdate_range("2026-01-05", periods=3) data = pd.DataFrame({"A": [10.0, 10.0, 10.0], "B": [10.0, 10.0, 10.0]}, index=idx) held = pd.DataFrame({"A": [0, 0.004, 0.0], "B": [0, 0.5, 0.5]}, index=idx) _, traded = simulate_live_book(held, data, 10_000.0, min_trade_value=50.0) self.assertEqual(traded["A"].abs().sum(), 0.0) # a $40 entry is below $50 self.assertGreater(traded["B"].abs().sum(), 0.0) class LotAllocationTests(unittest.TestCase): def test_leftover_cash_buys_the_lot_that_brings_a_name_nearest_its_target(self): prices = {"A": 45.0, "B": 13.0} # 70k book, 50/50: A wants 777.8 shares, B 2692.3. Rounding down leaves # 4.9k idle; it buys A's 8th lot (3.5k short of target vs half a lot # 2.25k), after which B's extra lot no longer fits. trades = trader.compute_trades({}, 70_000.0, {"A": 0.5, "B": 0.5}, prices, integer_shares=True, lot_size=100) shares = {t["ticker"]: t["shares_delta"] for t in trades} self.assertEqual(shares, {"A": 800, "B": 2600}) spent = sum(s * prices[t] for t, s in shares.items()) self.assertLessEqual(spent, 70_000.0) def test_a_lot_far_above_target_is_not_bought(self): trades = trader.compute_trades({}, 10_000.0, {"A": 0.3, "B": 0.7}, {"A": 70.0, "B": 1.0}, integer_shares=True, lot_size=100) # A's target is 3000 = 0.43 of a 7000 lot: stay out. self.assertNotIn("A", {t["ticker"] for t in trades}) def test_small_drift_is_not_traded_but_entries_and_exits_are(self): holdings = {"A": 10.0, "B": 10.0} prices = {"A": 105.0, "B": 95.0, "C": 50.0} trades = trader.compute_trades(holdings, 0.0, {"A": 0.5, "B": 0.5}, prices, min_trade_value=1.0, resize_band=0.10) self.assertEqual(trades, []) trades = trader.compute_trades(holdings, 0.0, {"A": 0.5, "C": 0.5}, prices, min_trade_value=1.0, resize_band=0.10) self.assertEqual({t["ticker"]: t["direction"] for t in trades}, {"B": "SELL", "C": "BUY"}) class PanelBackfillTests(unittest.TestCase): def test_tickers_missing_from_a_partial_batch_are_retried_before_giving_up(self): idx = pd.bdate_range("2026-01-05", periods=4) calls = [] def fake_download(tickers, start, end=None, fields=None): calls.append(list(tickers)) ok = [t for t in tickers if t == "AAA" or (t == "BBB" and len(calls) > 1)] frame = pd.DataFrame({t: [1.0, 2.0, 3.0, 4.0] for t in ok}, index=idx) return {f: frame for f in fields} with mock.patch.object(data_manager, "_download", fake_download): got, missing = data_manager._backfill(["AAA", "BBB", "CCC"], "2026-01-05", ["Close", "Volume"], idx) self.assertEqual(sorted(got["Close"].columns), ["AAA", "BBB"]) self.assertEqual(missing, ["CCC"]) self.assertEqual(calls[1], ["BBB", "CCC"]) class BacktestUniverseTests(unittest.TestCase): def test_stock_universe_excludes_benchmarks_and_etfs_sharing_the_panel(self): import backtest_runner idx = pd.bdate_range("2026-01-05", periods=3) panel = pd.DataFrame(1.0, index=idx, columns=["AAA", "BBB", "SPY", "TQQQ"]) fake_universe = {"us": {"fetch": lambda: ["AAA"], "benchmark": "SPY", "benchmark_label": "SPY"}} members = pd.DataFrame(True, index=idx, columns=["AAA", "BBB"]) with mock.patch.object(backtest_runner, "UNIVERSES", fake_universe), \ mock.patch.object(backtest_runner.uh, "pit_membership", lambda m, days: members.reindex(days).fillna(False)), \ mock.patch.object(backtest_runner, "_panels", lambda *a: {"close": panel, "open": panel}): data = backtest_runner.load_stock_panels("us", ["momentum"], None, refresh=False) self.assertEqual(data["stocks"], ["AAA", "BBB"]) class MembershipTests(unittest.TestCase): def test_csi300_snapshot_applies_until_the_next_one(self): snaps = {"2026-01-01": ["600000.SS", "000001.SZ"], "2026-02-01": ["600000.SS"]} dates = pd.to_datetime(["2025-12-31", "2026-01-15", "2026-02-02"]) mask = uh.csi300_membership_mask(dates, ["600000.SS", "000001.SZ", "510300.SS"], snaps) self.assertEqual(mask["000001.SZ"].tolist(), [False, True, False]) self.assertEqual(mask["600000.SS"].tolist(), [False, True, True]) # Never-member columns (benchmarks) are not masked. self.assertTrue(mask["510300.SS"].all()) def test_past_sessions_never_hold_a_name_before_it_joined_the_index(self): bars = _bars() bars["close"]["S00"] *= np.exp(np.linspace(0, 3, len(bars["close"]))) # the best momentum bars["open"]["S00"] = bars["close"]["S00"] days = bars["close"].index[-5:] others = [c for c in bars["close"].columns if c != "S00"] members = pd.DataFrame(True, index=bars["close"].index, columns=bars["close"].columns) members.loc[members.index < days[-1], "S00"] = False with mock.patch.object(trader, "index_members", lambda m, d: members.reindex(d)): weights = trader.session_signals("us", "momentum", others + ["S00"], days, bars["close"], bars["open"]) self.assertTrue((weights.loc[days[:-1], "S00"] == 0).all()) # The historical observations were public before index admission. self.assertGreater(weights.loc[days[-1], "S00"], 0) def test_changes_table_with_two_level_headers_is_found(self): cols = pd.MultiIndex.from_tuples([("Effective Date", "Effective Date"), ("Added", "Ticker"), ("Removed", "Ticker"), ("Reason", "Reason")]) table = pd.DataFrame([["September 21, 2026", "BE", "TAP", "cap"]], columns=cols) other = pd.DataFrame({"x": [1]}) changes = uh._find_changes_table([other, table]) self.assertEqual(list(changes.columns), ["Date", "Added_Ticker", "Removed_Ticker", "Reason_Reason"]) if __name__ == "__main__": unittest.main()