set50-system/backend/tests/test_backtest.py

"""Tests for the point-in-time multi-rebalance backtest engine."""

from __future__ import annotations

import datetime as dt
import unittest
from unittest.mock import patch

from app import backtest


def _fake_series() -> dict:
    """Synthetic daily price series for A (up) and B (flat), 300+ days."""
    def bars(base, drift):
        out = []
        for i in range(400):
            d = (dt.date(2025, 1, 1) + dt.timedelta(days=i)).isoformat()
            out.append({"date": d, "adjusted_close": base + drift * i})
        return out
    return {"A": {"bars": bars(10.0, 0.1)}, "B": {"bars": bars(20.0, 0.0)}}


def _flat_series() -> dict:
    return {"A": {"bars": [
        {"date": "2026-01-01", "adjusted_close": 10.0},
        {"date": "2026-02-01", "adjusted_close": 10.0},
    ]}}


def _rising_series() -> dict:
    return {"A": {"bars": [
        {"date": "2026-01-01", "adjusted_close": 10.0},
        {"date": "2026-02-01", "adjusted_close": 12.0},
    ]}}


class RebalanceDatesTest(unittest.TestCase):
    def test_monthly(self):
        d = backtest._rebalance_dates("2026-01-01", "2026-04-01")
        self.assertEqual(d, ["2026-01-01", "2026-02-01", "2026-03-01", "2026-04-01"])

    def test_quarterly_boundaries(self):
        d = backtest._rebalance_dates("2026-01-01", "2027-04-01", "quarterly")
        self.assertEqual(d, ["2026-01-01", "2026-04-01", "2026-07-01",
                             "2026-10-01", "2027-01-01", "2027-04-01"])

    def test_end_after_start_required(self):
        with self.assertRaises(backtest.BacktestError):
            backtest._rebalance_dates("2026-06-01", "2026-06-01")

    def test_bad_freq(self):
        with self.assertRaises(backtest.BacktestError):
            backtest._rebalance_dates("2026-01-01", "2026-06-01", "weekly")


class MomentumAtTest(unittest.TestCase):
    def test_true_12_1_skips_last_month(self):
        series = _fake_series()
        d = dt.date(2026, 5, 1)
        m = backtest.momentum_at(series, "A", d)
        # A rises 0.1/day; momentum over 12m should be clearly positive.
        self.assertIsNotNone(m)
        self.assertTrue(m is not None and m > 0.0)


class RunBacktestTest(unittest.TestCase):
    def _score_fn(self, syms, as_of):
        return {s: {"combined": 1.0 if s == "A" else 0.5,
                    "is_dividend": True, "dividend_yield": 2.0} for s in syms}

    @patch("app.backtest.load_price_snapshot", return_value=_fake_series())
    def test_multi_rebalance_reuses_portfolio(self, _load):
        res = backtest.run_backtest(
            "2026-01-01", "2026-06-01", capital=1_000_000,
            rebalance_freq="monthly", score_fn=self._score_fn,
            symbols=["A", "B"],
        )
        self.assertEqual(res.planned_rebalances, 6)
        # Real re-allocations happened (price data exists for every window).
        self.assertEqual(res.rebalances, 6)
        self.assertGreater(res.trades, 0)
        # Supplying a score_fn WITHOUT pit_meta is NOT PIT: the engine can no
        # longer trust an arbitrary callable. Only scores that assert
        # pit_meta.pit=True set leakage_guard (see the two tests below).
        self.assertFalse(res.leakage_guard)
        self.assertAlmostEqual(
            res.final_value,
            res.capital + res.price_pnl + res.dividend_income,
            places=2,
        )

    @patch("app.backtest.load_price_snapshot", return_value={})
    def test_no_price_snapshot_raises(self, _load):
        with self.assertRaises(backtest.BacktestError):
            backtest.run_backtest("2026-01-01", "2026-06-01")

    @patch("app.backtest.load_price_snapshot", return_value=_fake_series())
    def test_default_no_score_fn_non_pit(self, _load):
        # Without a score_fn, the engine uses the current board -> non-PIT.
        res = backtest.run_backtest(
            "2026-01-01", "2026-03-01", capital=100_000,
            rebalance_freq="monthly", symbols=["A", "B"],
        )
        self.assertFalse(res.leakage_guard)

    @patch("app.backtest.load_price_snapshot", return_value=_flat_series())
    def test_flat_price_has_zero_price_pnl(self, _load):
        def score_fn(symbols, as_of):
            return {"A": {"combined": 1.0, "is_dividend": False,
                          "dividend_yield": 0.0}}

        res = backtest.run_backtest(
            "2026-01-01", "2026-02-01", capital=100_000,
            score_fn=score_fn, symbols=["A"],
        )

        self.assertEqual(res.price_pnl, 0.0)
        self.assertEqual(res.dividend_income, 0.0)
        self.assertEqual(res.final_value, 100_000.0)
        self.assertEqual(res.net_return, 0.0)

    @patch("app.backtest.load_price_snapshot", return_value=_rising_series())
    def test_rising_price_without_dividend_is_all_price_pnl(self, _load):
        def score_fn(symbols, as_of):
            return {"A": {"combined": 1.0, "is_dividend": False,
                          "dividend_yield": 0.0}}

        res = backtest.run_backtest(
            "2026-01-01", "2026-02-01", capital=100_000,
            score_fn=score_fn, symbols=["A"],
        )

        # A non-dividend name is allocated through the canonical 20% bucket:
        # 20,000 invested at 10.0 rises 20%, producing 4,000 price P&L.
        self.assertEqual(res.price_pnl, 4_000.0)
        self.assertEqual(res.dividend_income, 0.0)
        self.assertEqual(res.final_value, 104_000.0)
        self.assertEqual(res.net_return, 0.04)

    @patch("app.backtest.load_price_snapshot", return_value=_flat_series())
    def test_dividend_is_included_in_final_value_and_net_return(self, _load):
        def score_fn(symbols, as_of):
            return {"A": {"combined": 1.0, "is_dividend": True,
                          "dividend_yield": 2.0}}

        res = backtest.run_backtest(
            "2026-01-01", "2026-02-01", capital=100_000,
            score_fn=score_fn, symbols=["A"],
        )

        self.assertEqual(res.price_pnl, 0.0)
        self.assertEqual(res.dividend_income, 1_000.0)
        self.assertEqual(res.final_value, 101_000.0)
        self.assertEqual(res.net_return, 0.01)
        self.assertEqual(
            res.to_dict()["dividend_method"],
            "final_holdings_yield_proxy",
        )
        self.assertEqual(
            res.final_value,
            res.capital + res.price_pnl + res.dividend_income,
        )

    @patch("app.backtest.load_price_snapshot", return_value=_fake_series())
    def test_supplied_fn_without_pit_meta_is_not_pit(self, _load):
        # A supplied score_fn that does NOT assert PIT integrity via pit_meta
        # must NOT set leakage_guard (the old behaviour trusted any callable).
        def naive(syms, as_of=None):
            return {s: {"combined": 0.5, "is_dividend": True, "dividend_yield": 5.0}
                    for s in syms}
        res = backtest.run_backtest(
            "2026-01-01", "2026-03-01", capital=100_000,
            rebalance_freq="monthly", score_fn=naive, symbols=["A", "B"],
        )
        self.assertIs(res.leakage_guard, False)

    @patch("app.backtest.load_price_snapshot", return_value=_fake_series())
    def test_supplied_fn_with_pit_meta_sets_leakage_guard(self, _load):
        # Only a score_fn whose scores assert pit_meta.pit=True may set guard.
        def pit(syms, as_of=None):
            return {
                s: {"combined": 0.5, "is_dividend": True, "dividend_yield": 5.0,
                     "pit_meta": {"pit": True, "partial_pit": True, "note": "pit"}}
                for s in syms
            }
        res = backtest.run_backtest(
            "2026-01-01", "2026-03-01", capital=100_000,
            rebalance_freq="monthly", score_fn=pit, symbols=["A", "B"],
        )
        self.assertIs(res.leakage_guard, True)


if __name__ == "__main__":
    unittest.main()