"""Unit tests for the after-tax model (tax.py) — synthetic, no data bundle. Key regression: distributions are RAW per-share dollars, so the account must hold RAW share units (valued on close prices). The old code valued units on ADJ (total-return index) prices, which overstates every distribution — and its tax — by the raw/adj ratio (e.g. a fund that has paid many distributions: adj 11 vs raw 30 -> 2.7x overcharge). Run: .venv/bin/python tests/test_tax.py """ from __future__ import annotations import sys from pathlib import Path import pandas as pd sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import tax # noqa: E402 PASS, FAIL = 0, 0 def check(name: str, ok: bool) -> None: global PASS, FAIL PASS += bool(ok) FAIL += not ok print((" ok " if ok else " FAIL ") + name) def main() -> int: # One fund, raw prices 30 -> 34, ex-div 5.00 on day 6 (34 -> 29), # then 29.5. Account starts at 1.0 = 1/30 share. dates = pd.to_datetime(["2024-01-0%d" % i for i in range(1, 8)]) close = pd.DataFrame({"f": [30.0, 31.0, 32.0, 33.0, 34.0, 29.0, 29.5]}, index=dates) div = pd.DataFrame({"f": [0, 0, 0, 0, 0, 0.0, 0.0]}, index=dates) capg = pd.DataFrame({"f": [0, 0, 0, 0, 0, 5.0, 0.0]}, index=dates) r = tax.after_tax_portfolio(close, div, capg, {"f": 1.0}, lt_rate=0.20, st_rate=0.15, div_rate=0.15) t5 = dates[5] # --- the regression: tax on the distribution is 20% of 5/30 of the # account (the distribution as a fraction of the RAW price), i.e. # 3.333% — not 5/adj (which the old code implied, ~4% with this # adj path and up to 2.7x for high-history funds). check("cap-gain tax is 20% of (5/30) of the account", abs(r.taxes["capg_tax"][t5] - 0.2 * 5.0 / 30.0) < 1e-12) check("no dividend tax", r.taxes["div_tax"].abs().sum() == 0.0) # --- reinvestment: the after-tax remainder (5*0.8 of the 1/30 share) # is bought at the ex-div close 29, so MV on the event day is # (1/30)*(29 + 4) = 33/30. The tax is already inside the (smaller) # units — it must NOT also show up in equity as a second deduction. mv5 = r.equity[t5] + r.liq_tax[t5] # no cash component on this day check("market value on ex-div day = 33/30 (after-tax reinvested)", abs(mv5 - 33.0 / 30.0) < 1e-9) # --- liquidation value: price drop (5/30) is exactly offset by the # distribution kept, so equity falls by the true tax cost only. # Day-5 equity = MV(33/30) - liq_tax(0: basis 1+4/30 > MV, a "loss" # under the after-tax basis rule). check("equity on ex-div day = 33/30 (tax cost lives in the units)", abs(r.equity[t5] - 33.0 / 30.0) < 1e-9) # --- pre-event day: plain ST-marked liquidation value. t4 = dates[4] check("equity day before = MV - 15% ST gain tax", abs(r.equity[t4] - (34.0 / 30.0) * (1 - 0.15 * (4.0 / 34.0))) < 1e-9) # --- cross-day consistency: equity(5)/equity(4) from the explicit # values (day 4 = 34/30 MV less 15% ST gain tax; day 5 = 33/30). check("ex-div equity step matches explicit values", abs(r.equity[t5] / r.equity[t4] - (33.0 / 30.0) / (34.0 / 30.0 - 0.15 * 4.0 / 30.0)) < 1e-9) # --- no rebalancing: one lot survives, units grew by the reinvested # net amount at the ex-div price. check("single lot outstanding", r.lots_outstanding == 1) # --- dividend leg taxed at div_rate, cap-gain leg at lt_rate: # split a 1.00 distribution 0.4 div / 0.6 capg on day 6. div2 = pd.DataFrame({"f": [0, 0, 0, 0, 0, 0.4, 0.0]}, index=dates) capg2 = pd.DataFrame({"f": [0, 0, 0, 0, 0, 0.6, 0.0]}, index=dates) r2 = tax.after_tax_portfolio(close, div2, capg2, {"f": 1.0}, lt_rate=0.20, st_rate=0.15, div_rate=0.10) check("split distribution taxed by component rate", abs(r2.taxes["capg_tax"][t5] - 0.20 * 0.6 / 30.0) < 1e-12 and abs(r2.taxes["div_tax"][t5] - 0.10 * 0.4 / 30.0) < 1e-12) # --- state/local + NIIT: composite rates add to each component # (NY/NYC tax cap gains at ordinary rates — the caller passes the # composite sl_rate; the model must not preferentiate cap gains). r3 = tax.after_tax_portfolio(close, div2, capg2, {"f": 1.0}, lt_rate=0.20, st_rate=0.15, div_rate=0.10, niit=0.038, sl_rate=0.14) check("capg dist taxed at lt + NIIT + state/local", abs(r3.taxes["capg_tax"][t5] - (0.20 + 0.038 + 0.14) * 0.6 / 30.0) < 1e-12) check("div taxed at div + NIIT + state/local", abs(r3.taxes["div_tax"][t5] - (0.10 + 0.038 + 0.14) * 0.4 / 30.0) < 1e-12) # realized-gain leg: two symbols, daily rebalance forces ST sells; # every realized dollar is taxed at (st_rate + niit + sl_rate). d3 = pd.to_datetime(["2024-01-0%d" % i for i in (1, 2, 3)]) close3 = pd.DataFrame({"a": [30.0, 34.0, 29.0], "b": [10.0, 10.0, 10.0]}, index=d3) zero3 = pd.DataFrame(0.0, index=d3, columns=["a", "b"]) r4 = tax.after_tax_portfolio(close3, zero3, zero3, {"a": 0.5, "b": 0.5}, rebalance="1D", lt_rate=0.20, st_rate=0.15, div_rate=0.10, niit=0.038, sl_rate=0.14) real_tax = float(r4.taxes["realized_tax"].sum()) net_st = float(r4.realized["st_gain"].sum()) - \ float(r4.realized["st_loss"].sum()) check("realized ST gains/losses net-taxed at st + NIIT + state/local", float(r4.realized["st_gain"].sum()) > 0 and abs(real_tax - net_st * (0.15 + 0.038 + 0.14)) < 1e-9) # --- EQUITY PATH (not just recorded taxes) with the full stack: # the reinvested net must deduct (federal + NIIT + state), and the # liquidation tax on each date must use the same stacked rates. # A: cap-gain dist 5.0 on a 30-priced fund; keep = 1-0.20-0.038-0.14 zero_d = pd.DataFrame(0.0, index=dates, columns=["f"]) r5 = tax.after_tax_portfolio(close, zero_d, capg, {"f": 1.0}, lt_rate=0.20, st_rate=0.15, div_rate=0.10, niit=0.038, sl_rate=0.14) keep = 1.0 - 0.20 - 0.038 - 0.14 mv5 = float(r5.equity[t5]) + float(r5.liq_tax[t5]) check("reinvested net deducts NIIT + state (MV = 29/30 + net/30)", abs(mv5 - (29.0 / 30.0 + (5.0 / 30.0) * keep)) < 1e-9) # B: sale tax on each calculation date — ST window (2 days) dlt = pd.to_datetime(["2024-01-01", "2024-01-02"]) c2 = pd.DataFrame({"f": [30.0, 34.0]}, index=dlt) r6 = tax.after_tax_portfolio(c2, pd.DataFrame(0.0, index=dlt, columns=["f"]), pd.DataFrame(0.0, index=dlt, columns=["f"]), {"f": 1.0}, lt_rate=0.20, st_rate=0.15, div_rate=0.10, niit=0.038, sl_rate=0.14) check("liquidation tax (ST) uses st + NIIT + state", abs(float(r6.equity[dlt[1]]) - (34.0 / 30.0 - (0.15 + 0.038 + 0.14) * 4.0 / 30.0)) < 1e-9) # B2: ...and the LT leg dlt2 = pd.to_datetime(["2022-01-03", "2024-01-02"]) c2b = pd.DataFrame({"f": [30.0, 34.0]}, index=dlt2) r7 = tax.after_tax_portfolio(c2b, pd.DataFrame(0.0, index=dlt2, columns=["f"]), pd.DataFrame(0.0, index=dlt2, columns=["f"]), {"f": 1.0}, lt_rate=0.20, st_rate=0.15, div_rate=0.10, niit=0.038, sl_rate=0.14) check("liquidation tax (LT) uses lt + NIIT + state", abs(float(r7.equity[dlt2[1]]) - (34.0 / 30.0 - (0.20 + 0.038 + 0.14) * 4.0 / 30.0)) < 1e-9) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 if __name__ == "__main__": sys.exit(main())