f/tests/test_data.py

"""Data-cache tests: incremental refresh from changed data files.

Uses a small synthetic data root in a temp dir — no real data, no server.
Run:  .venv/bin/python tests/test_data.py
"""
from __future__ import annotations

import json
import os
import pathlib
import shutil
import sys
import tempfile
import time

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

import data as D  # noqa: E402

PASS, FAIL = 0, 0


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 write_history(root: pathlib.Path, sym: str, last_price: float | None = None,
                  days: int = 10) -> pathlib.Path:
    """Prices 100.00, 101.00, ..., 109.00 (last day overridable)."""
    f = root / f"{sym}-history.csv"
    lines = ["Date,Open,High,Low,Close,Adj Close,Volume"]
    for i in range(days):
        p = last_price if (last_price is not None and i == days - 1) \
            else 100.0 + i
        lines.append(f"2024-01-{i + 1:02d},{p:.2f},{p:.2f},{p:.2f},{p:.2f},{p:.2f},1000")
    f.write_text("\n".join(lines) + "\n")
    return f


def touch(p: pathlib.Path) -> None:
    """Bump mtime so the manifest sees the file as changed."""
    st = p.stat()
    os.utime(p, ns=(st.st_atime_ns + 10**9, st.st_mtime_ns + 10**9))


def reload(d: pathlib.Path) -> D.Bundle:
    # refresh() is the synchronous path: deterministic for these assertions
    return D.refresh(root=d["root"], cache=d["cache"])


def main() -> int:
    # fast watcher for the background-refresh test below
    D.WATCH_INTERVAL, D.REFRESH_CADENCE = 0.2, 0.0
    # synthetic root: no frozen overrides
    tmp = pathlib.Path(tempfile.mkdtemp(prefix="fdatatest-"))
    D.FROZEN_DIR = tmp / "frozen"
    D.FROZEN_DIR.mkdir()
    root, cache = tmp / "stocks", tmp / "cache"
    root.mkdir()
    d = {"root": root, "cache": cache}
    try:
        write_history(root, "aaa")
        write_history(root, "bbb")
        write_history(root, "ccc")
        (root / "aaa-dividend.csv").write_text("Date,Dividends\n2024-01-05,0.5\n")
        (root / "bbb.json").write_text(
            json.dumps({"chart": {"result": [{"meta": {"longName": "Bee Bee Corp"}}]}}))

        b = reload(d)
        check("initial full build", list(b.adj.columns) == ["aaa", "bbb", "ccc"]
              and b.adj["aaa"].iloc[-1] == 109.0, str(list(b.adj.columns)))
        check("names from json", b.names.get("bbb") == "Bee Bee Corp")
        check("dividend panel", b.div["aaa"].loc["2024-01-05"] == 0.5)

        # --- modified file: last price changes
        f = write_history(root, "aaa", last_price=200.0)
        touch(f)
        b = reload(d)
        check("changed file re-read incrementally", b.adj["aaa"].iloc[-1] == 200.0)
        check("untouched symbols intact", b.adj["bbb"].iloc[-1] == 109.0)

        # --- new file: symbol added
        write_history(root, "ddd")
        b = reload(d)
        check("new symbol added", "ddd" in b.adj.columns)

        # --- deleted file: symbol dropped
        (root / "ccc-history.csv").unlink()
        b = reload(d)
        check("deleted symbol dropped", "ccc" not in b.adj.columns
              and "ccc" not in b.div.columns)

        # --- name file added for a new symbol
        (root / "ddd.json").write_text(
            json.dumps({"chart": {"result": [{"meta": {"longName": "Dee Dee"}}]}}))
        b = reload(d)
        check("name added incrementally", b.names.get("ddd") == "Dee Dee")

        # --- dividend file modified
        (root / "aaa-dividend.csv").write_text(
            "Date,Dividends\n2024-01-05,0.5\n2024-01-06,1.25\n")
        b = reload(d)
        check("dividend change picked up", b.div["aaa"].loc["2024-01-06"] == 1.25)

        # --- no-change reload: cheap, no rewrite
        before = (cache / "panel_adj.parquet").stat().st_mtime_ns
        b = reload(d)
        after = (cache / "panel_adj.parquet").stat().st_mtime_ns
        check("no-op reload does not rewrite parquet", before == after)

        # --- background refresh: load_bundle never blocks on changed data;
        #     it serves the current bundle and a background thread catches up
        f = write_history(root, "aaa", last_price=300.0)
        touch(f)
        t0 = time.monotonic()
        b = D.load_bundle(root=root, cache=cache)
        check("load_bundle returns immediately on changed data",
              time.monotonic() - t0 < 2.0)
        deadline = time.monotonic() + 15.0
        while time.monotonic() < deadline:
            b = D.load_bundle(root=root, cache=cache)
            if b.adj["aaa"].iloc[-1] == 300.0:
                break
            time.sleep(0.2)
        check("background thread picks up the change", b.adj["aaa"].iloc[-1] == 300.0)
        check("generation counter advances", D.generation(root=root, cache=cache) >= 1)

        # --- forced full rebuild still works
        b = D.load_bundle(root=root, cache=cache, rebuild=True)
        check("forced full rebuild", list(b.adj.columns) == ["aaa", "bbb", "ddd"])
    finally:
        shutil.rmtree(tmp, ignore_errors=True)
    print(f"\n{PASS} passed, {FAIL} failed")
    return 1 if FAIL else 0


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