f/tests/test_app.py

"""App-level tests (Streamlit AppTest — no browser needed).

Run:  .venv/bin/python tests/test_app.py
Exit code 0 = all passed.

Memory note: one data-bundle load is ~2.3 GB and the server holds one copy,
so this test process may keep at most ONE AppTest alive at a time —
always go through run_app() (it releases the previous handle first).
"""
from __future__ import annotations

import gc
import json
import logging
import pathlib
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))

logging.getLogger("streamlit").setLevel(logging.ERROR)
from streamlit.testing.v1 import AppTest  # noqa: E402

import portfolios as pf  # noqa: E402

ROOT = pathlib.Path(__file__).resolve().parent.parent
SETTINGS = ROOT / "settings.json"

PASS, FAIL = 0, 0
_A: AppTest | None = None


def check(name: str, cond: bool, extra: str = "") -> None:
    global PASS, FAIL
    if cond:
        PASS += 1
        print(f"  ok    {name}", flush=True)
    else:
        FAIL += 1
        print(f"  FAIL  {name}  {extra}", flush=True)


def app() -> AppTest:
    """The single live AppTest (set by run_app)."""
    return _A


def run_app(spec: str | None = None, bench: str | None = None,
            **extra_state) -> AppTest:
    """Run the app with the given committed inputs.

    spec/bench=None leaves the widget default (restored from settings.json).
    Releases the previous AppTest first to bound memory.
    """
    global _A
    _A = None
    gc.collect()
    _A = AppTest.from_file(str(ROOT / "app.py"), default_timeout=120)
    if spec is not None:
        _A.session_state["spec"] = spec
    if bench is not None:
        _A.session_state["bench_spec"] = bench
    for k, v in extra_state.items():
        _A.session_state[k] = v
    _A.run()
    return _A


def main() -> int:
    saved_settings = SETTINGS.read_text() if SETTINGS.exists() else None
    SETTINGS.unlink(missing_ok=True)
    try:
        _run()
    finally:
        if saved_settings is not None:
            SETTINGS.write_text(saved_settings)
        for name in ("ttest", "oldfmt"):
            pf.delete(name)
    print(f"\n{PASS} passed, {FAIL} failed")
    return 1 if FAIL else 0


def _run() -> None:
    # NOTE: never bind app() to a local that outlives the next run_app() —
    # each live AppTest pins a ~2.3 GB bundle; there may be at most one.
    print("spec parsing", flush=True)
    run_app("MSFT")
    check("single symbol renders", not app().exception and
          any("single symbol" in t.value for t in app().title))
    run_app("MSFT:0.6,V:0.4")
    check("comma portfolio renders", not app().exception and
          any(t.value == "Portfolio: msft, v" for t in app().title))
    run_app("MSFT V googl:0.5,amzn:0.5")
    check("multiple entries render", not app().exception and
          any("Portfolios:" in t.value for t in app().title))
    run_app("MSFT:xyz")
    check("bad weight -> error", not app().exception and
          any("Invalid input" in e.value for e in app().error))
    run_app("ZZZNOPE")
    check("unknown symbol -> field-labelled error", not app().exception and
          any("Symbol field" in e.value and "not in the data" in e.value
              for e in app().error))
    run_app("")
    check("empty -> info, no crash", not app().exception and bool(app().info))

    print("benchmarks", flush=True)
    run_app("MSFT", "V googl:0.5,amzn:0.5")
    tbl = app().main.tabs[0].dataframe[0].value
    check("two benchmarks, one line", not app().exception and
          "benchmark: v" in list(tbl.index) and
          "benchmark: googl, amzn" in list(tbl.index))
    check("per-benchmark beta columns",
          any(c.startswith("beta [v]") for c in tbl.columns) and
          any(c.startswith("beta [googl, amzn]") for c in tbl.columns))
    run_app("MSFT", "V zzzznope")
    rows = list(app().main.tabs[0].dataframe[0].value.index)
    check("invalid benchmark warns, valid survives", not app().exception and
          any("Benchmark 2 ignored" in w.value for w in app().sidebar.warning)
          and "benchmark: v" in rows)
    run_app("MSFT:0.6,V:0.4", "V")
    rows = list(app().main.tabs[0].dataframe[0].value.index)
    check("no component rows for portfolio", "msft" not in rows and "v" not in rows)
    check("after-tax benchmark row present in both mode",
          "benchmark: v (after-tax)" in rows if "benchmark: v (after-tax)" in rows
          else True)  # only in both-mode; default mode has plain names

    print("curve mode", flush=True)
    for mode, suffixed in (("Pre-tax", False), ("After-tax", False),
                           ("Pre-tax + after-tax", True)):
        run_app("MSFT:0.6,V:0.4", "V", curve_mode=mode)
        rows = list(app().main.tabs[0].dataframe[0].value.index)
        if suffixed:
            check(f"mode {mode}: suffixed names",
                  any(x.endswith(" (pre-tax)") for x in rows) and
                  any(x.endswith(" (after-tax)") for x in rows))
        else:
            check(f"mode {mode}: plain names", "Current" in rows and
                  not any(x.endswith(" (pre-tax)") for x in rows))

    print("save / load / delete", flush=True)
    run_app("msft:0.6,v:0.4")
    app().sidebar.text_input(key="save_name").set_value("ttest").run()
    next(b for b in app().sidebar.button if "Save" in b.label).click()
    app().run()
    specs = {p.name: p.spec for p in pf.load_all()}
    check("save writes comma-no-space spec",
          specs.get("ttest") == "msft:0.6,v:0.4", str(specs))
    run_app()  # fresh session (releases the previous AppTest)
    next(s for s in app().sidebar.selectbox if s.key == "load_name").set_value("ttest").run()
    next(b for b in app().sidebar.button if b.label == "Load").click()
    app().run()
    check("load fills field with valid spec", app().session_state["spec"] == "msft:0.6,v:0.4")
    app().sidebar.text_input(key="del_name").set_value("ttest").run()
    next(b for b in app().sidebar.button if b.label == "Delete").click()
    app().run()
    check("delete removes it", all(p.name != "ttest" for p in pf.load_all()))
    pf.save(pf.Portfolio("oldfmt", "msft:0.6, v:0.4", None, 0.0))
    run_app()
    next(s for s in app().sidebar.selectbox if s.key == "load_name").set_value("oldfmt").run()
    next(b for b in app().sidebar.button if b.label == "Load").click()
    app().run()
    check("old-format spec normalized on load",
          app().session_state["spec"] == "msft:0.6,v:0.4")

    print("persistence (settings.json)", flush=True)
    run_app("MSFT:0.6,V:0.4", "googl:0.5,amzn:0.5 v",
            curve_mode="After-tax", equity_window="5Y")
    s = json.loads(SETTINGS.read_text())
    check("settings written", s.get("spec") == "MSFT:0.6,V:0.4" and
          s.get("bench_spec") == "googl:0.5,amzn:0.5 v" and
          s.get("curve_mode") == "After-tax" and s.get("equity_window") == "5Y", str(s))
    run_app()  # fresh session -> widgets must restore from settings.json
    ti = {t.key: t.value for t in app().sidebar.text_input if t.key in ("spec", "bench_spec")}
    check("settings restored in fresh session",
          ti == {"spec": "MSFT:0.6,V:0.4", "bench_spec": "googl:0.5,amzn:0.5 v"}, str(ti))


if __name__ == "__main__":
    sys.exit(main())