60 lines
2.2 KiB
Python
60 lines
2.2 KiB
Python
|
"""Tests for the flex-spending engine."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from fire_planner.flex_spending import FlexRule, applicable_cut, cuts_per_year
|
||
|
|
|
||
|
|
|
||
|
|
def test_applicable_cut_picks_deepest_rule() -> None:
|
||
|
|
rules = [
|
||
|
|
FlexRule(from_ath_pct=0.10, cut_discretionary_pct=0.20),
|
||
|
|
FlexRule(from_ath_pct=0.30, cut_discretionary_pct=0.60),
|
||
|
|
FlexRule(from_ath_pct=0.50, cut_discretionary_pct=0.90),
|
||
|
|
]
|
||
|
|
# No drawdown — no cut.
|
||
|
|
assert applicable_cut(0.0, rules) == 0.0
|
||
|
|
# 9% drop — below first threshold.
|
||
|
|
assert applicable_cut(0.09, rules) == 0.0
|
||
|
|
# 15% drop — only first rule fires.
|
||
|
|
assert applicable_cut(0.15, rules) == pytest.approx(0.20)
|
||
|
|
# 35% drop — first + second; deepest cut wins (0.60, not 0.80).
|
||
|
|
assert applicable_cut(0.35, rules) == pytest.approx(0.60)
|
||
|
|
# 60% drop — all three; 0.90 wins.
|
||
|
|
assert applicable_cut(0.60, rules) == pytest.approx(0.90)
|
||
|
|
|
||
|
|
|
||
|
|
def test_applicable_cut_empty_rules() -> None:
|
||
|
|
assert applicable_cut(0.5, []) == 0.0
|
||
|
|
|
||
|
|
|
||
|
|
def test_cuts_per_year_handles_running_ath() -> None:
|
||
|
|
# Single path. Year 0 seed=1000, year 1 = 1200 (new ATH), year 2 = 800
|
||
|
|
# (-33% from ATH 1200), year 3 = 900 (still -25% from ATH 1200), year
|
||
|
|
# 4 = 1300 (new ATH).
|
||
|
|
portfolio = np.array([[1000, 1200, 800, 900, 1300]], dtype=np.float64)
|
||
|
|
rules = [
|
||
|
|
FlexRule(from_ath_pct=0.10, cut_discretionary_pct=0.20),
|
||
|
|
FlexRule(from_ath_pct=0.30, cut_discretionary_pct=0.60),
|
||
|
|
]
|
||
|
|
cuts = cuts_per_year(portfolio, rules)
|
||
|
|
# cuts[:, y] uses portfolio[:, y] (start-of-year decision based on
|
||
|
|
# the prior year's close).
|
||
|
|
# y=0: portfolio=1000 == ATH → 0
|
||
|
|
# y=1: portfolio=1200 == ATH → 0
|
||
|
|
# y=2: drawdown = 1 - 800/1200 = 0.333 → 0.60
|
||
|
|
# y=3: drawdown = 1 - 900/1200 = 0.25 → 0.20
|
||
|
|
assert cuts.shape == (1, 4)
|
||
|
|
assert cuts[0, 0] == pytest.approx(0.0)
|
||
|
|
assert cuts[0, 1] == pytest.approx(0.0)
|
||
|
|
assert cuts[0, 2] == pytest.approx(0.60)
|
||
|
|
assert cuts[0, 3] == pytest.approx(0.20)
|
||
|
|
|
||
|
|
|
||
|
|
def test_cuts_per_year_no_rules_returns_zeros() -> None:
|
||
|
|
portfolio = np.array([[1000, 800, 600]], dtype=np.float64)
|
||
|
|
cuts = cuts_per_year(portfolio, [])
|
||
|
|
assert cuts.shape == (1, 2)
|
||
|
|
assert (cuts == 0).all()
|