fire-planner/tests/test_spending_plan.py

"""SpendingPlanStrategy: user-customisable initial spend + annual adjust + guardrail."""
from fire_planner.strategies.base import StrategyState
from fire_planner.strategies.spending_plan import SpendingPlanStrategy


def state(**overrides: float | int) -> StrategyState:
    base = dict(
        portfolio=1_000_000.0,
        initial_portfolio=1_000_000.0,
        initial_withdrawal=40_000.0,
        year_idx=0,
        horizon_years=60,
        last_withdrawal=40_000.0,
        expected_real_return=0.04,
    )
    base.update(overrides)
    return StrategyState(**base)  # type: ignore[arg-type]


def test_year_zero_takes_initial_withdrawal() -> None:
    s = SpendingPlanStrategy()
    assert s.propose_withdrawal(state()) == 40_000.0


def test_year_zero_explicit_override() -> None:
    s = SpendingPlanStrategy(initial_spend=72_000.0)
    assert s.propose_withdrawal(state()) == 72_000.0


def test_constant_real_with_zero_adjust() -> None:
    s = SpendingPlanStrategy(annual_real_adjust_pct=0.0)
    assert s.propose_withdrawal(state(year_idx=10, last_withdrawal=40_000.0)) == 40_000.0


def test_positive_annual_adjust_grows_spending() -> None:
    """+2% real adjust → year 5 spend is last_w × 1.02 (one step from year 4)."""
    s = SpendingPlanStrategy(annual_real_adjust_pct=0.02)
    out = s.propose_withdrawal(state(year_idx=5, last_withdrawal=40_000.0))
    assert abs(out - 40_800.0) < 1e-6


def test_negative_annual_adjust_shrinks_spending() -> None:
    s = SpendingPlanStrategy(annual_real_adjust_pct=-0.005)
    out = s.propose_withdrawal(state(year_idx=10, last_withdrawal=40_000.0))
    assert abs(out - 39_800.0) < 1e-6


def test_guardrail_not_triggered_when_portfolio_above_threshold() -> None:
    s = SpendingPlanStrategy(guardrail_threshold_pct=0.80, guardrail_cut_pct=0.10)
    # portfolio at 90% of starting → above threshold → no cut
    out = s.propose_withdrawal(state(year_idx=3, portfolio=900_000.0,
                                     last_withdrawal=40_000.0))
    assert out == 40_000.0


def test_guardrail_triggered_cuts_by_pct() -> None:
    s = SpendingPlanStrategy(guardrail_threshold_pct=0.80, guardrail_cut_pct=0.10)
    # portfolio at 70% of starting → below threshold → 10% cut
    out = s.propose_withdrawal(state(year_idx=3, portfolio=700_000.0,
                                     last_withdrawal=40_000.0))
    assert abs(out - 36_000.0) < 1e-6


def test_guardrail_combines_with_annual_adjust() -> None:
    """Adjust applies first, then cut — both fire."""
    s = SpendingPlanStrategy(
        annual_real_adjust_pct=0.02,
        guardrail_threshold_pct=0.80,
        guardrail_cut_pct=0.10,
    )
    out = s.propose_withdrawal(state(year_idx=3, portfolio=700_000.0,
                                     last_withdrawal=40_000.0))
    # 40_000 * 1.02 = 40_800; trigger; 40_800 * 0.90 = 36_720
    assert abs(out - 36_720.0) < 1e-6


def test_guardrail_disabled_when_threshold_none() -> None:
    s = SpendingPlanStrategy(guardrail_threshold_pct=None)
    out = s.propose_withdrawal(state(year_idx=3, portfolio=10_000.0,
                                     last_withdrawal=40_000.0))
    assert out == 40_000.0  # no cut despite tiny portfolio


def test_returns_zero_when_portfolio_drained() -> None:
    s = SpendingPlanStrategy()
    out = s.propose_withdrawal(state(year_idx=5, portfolio=0.0, last_withdrawal=40_000.0))
    assert out == 0.0