easy_tdx_max/tests/unit/test_task_store.py

"""回测任务 SQLite 持久化测试（task_store + task_runner 集成 + REST 导出）。

覆盖：
- ``TaskStore``：save/load/list_recent/delete 往返、淘汰、重启恢复中断任务
- ``BacktestTaskRunner`` 集成：任务完成后落盘、内存淘汰后磁盘兜底、
  「重启」（新建 runner）后仍可查历史任务
- REST 导出端点：JSON 全量 / CSV 主表 / 未完成任务拒绝导出

持久化默认被 ``tests/conftest.py`` 关闭（``EASY_TDX_NO_TASK_DB=1``），
本文件的测试显式删除该变量并把 ``EASY_TDX_CONFIG_DIR`` 指向 ``tmp_path``。
"""

from __future__ import annotations

import time

import pytest

pytest.importorskip("fastapi")

from fastapi.testclient import TestClient  # noqa: E402

from easy_tdx.web import task_store as ts_mod  # noqa: E402
from easy_tdx.web.task_runner import BacktestTaskRunner  # noqa: E402
from easy_tdx.web.task_store import TaskStore  # noqa: E402


@pytest.fixture()
def persisted_env(tmp_path, monkeypatch):
    """开启持久化并指向临时目录；隔离全局单例。"""
    monkeypatch.delenv("EASY_TDX_NO_TASK_DB", raising=False)
    monkeypatch.setenv("EASY_TDX_CONFIG_DIR", str(tmp_path))
    ts_mod.reset_task_store()
    yield tmp_path
    ts_mod.reset_task_store()


def _wait_done(runner: BacktestTaskRunner, task_id: str, timeout: float = 5.0) -> None:
    """轮询等待任务进入终态。"""
    deadline = time.time() + timeout
    while time.time() < deadline:
        state = runner.peek(task_id)
        assert state is not None
        if state.status in ("done", "failed"):
            return
        time.sleep(0.02)
    raise AssertionError(f"任务 {task_id} 超时未完成")


# ── TaskStore 单元 ────────────────────────────────────────────────────────────


def test_store_save_load_roundtrip(persisted_env):
    store = TaskStore(db_path=persisted_env / "t.db")
    store.save(
        task_id="abc",
        status="done",
        description="ma_cross | 300根",
        created_at=1000.0,
        started_at=1001.0,
        finished_at=1002.0,
        result={"performance": {"total_return": 0.25}, "trades": [{"pnl": 1.0}]},
    )
    d = store.load("abc")
    assert d is not None
    assert d["status"] == "done"
    assert d["result"]["performance"]["total_return"] == 0.25
    assert d["created_at"] == 1000.0
    assert store.load("missing") is None


def test_store_list_recent_order_and_delete(persisted_env):
    store = TaskStore(db_path=persisted_env / "t.db")
    for i in range(5):
        store.save(task_id=f"t{i}", status="done", created_at=1000.0 + i)
    rows = store.list_recent(limit=3)
    assert [r["task_id"] for r in rows] == ["t4", "t3", "t2"]  # created_at 倒序
    assert store.delete("t4") is True
    assert store.delete("t4") is False
    assert store.load("t4") is None


def test_store_upsert_replaces(persisted_env):
    """同 task_id 二次 save 是覆盖而非追加。"""
    store = TaskStore(db_path=persisted_env / "t.db")
    store.save(task_id="x", status="running", created_at=1.0)
    store.save(task_id="x", status="done", created_at=1.0, finished_at=2.0, result={"a": 1})
    d = store.load("x")
    assert d["status"] == "done"
    assert d["result"] == {"a": 1}
    assert len(store.list_recent(limit=10)) == 1


def test_store_recovers_interrupted_tasks(persisted_env):
    """新连接（模拟进程重启）把遗留 pending/running 标记为 failed。"""
    store = TaskStore(db_path=persisted_env / "t.db")
    store.save(task_id="p1", status="pending", created_at=1.0)
    store.save(task_id="r1", status="running", created_at=1.0)
    store.save(task_id="d1", status="done", created_at=1.0, result={"ok": True})

    # 模拟重启：新实例初始化时触发恢复
    store2 = TaskStore(db_path=persisted_env / "t.db")
    _ = store2.list_recent(limit=10)

    assert store2.load("p1")["status"] == "failed"
    assert "重启" in store2.load("p1")["error"]
    assert store2.load("r1")["status"] == "failed"
    assert store2.load("d1")["status"] == "done"  # 已完成任务不受影响


def test_store_result_json_corruption_degrades(persisted_env):
    """result_json 损坏时 load 降级返回 result=None，不抛异常。"""
    import sqlite3

    path = persisted_env / "t.db"
    store = TaskStore(db_path=path)
    store.save(task_id="bad", status="done", created_at=1.0, result={"a": 1})
    conn = sqlite3.connect(path)
    conn.execute("UPDATE backtest_tasks SET result_json = '{not-json' WHERE task_id='bad'")
    conn.commit()
    conn.close()

    store2 = TaskStore(db_path=path)
    d = store2.load("bad")
    assert d is not None
    assert d["result"] is None


# ── Runner 集成 ────────────────────────────────────────────────────────────────


def test_runner_persists_done_task_and_survives_memory_eviction(persisted_env):
    runner = BacktestTaskRunner(max_workers=2, max_results=2)
    task_id = runner.submit(lambda: {"performance": {"total_return": 0.5}}, description="d")
    _wait_done(runner, task_id)

    # 磁盘上能查到 done + 完整结果
    d = ts_mod.get_task_store().load(task_id)
    assert d is not None
    assert d["status"] == "done"
    assert d["result"]["performance"]["total_return"] == 0.5

    # 内存淘汰（提交 3 个新任务挤掉 LRU）后 peek 仍能从磁盘兜底
    for _ in range(3):
        _wait_done(runner, runner.submit(lambda: {"x": 1}))
    state = runner.peek(task_id)
    assert state is not None
    assert state.status == "done"
    assert state.result["performance"]["total_return"] == 0.5
    runner.shutdown()


def test_runner_new_instance_sees_history(persisted_env):
    """「重启」：全新 runner/store 仍能列出并查询历史任务。"""
    runner1 = BacktestTaskRunner(max_workers=1)
    task_id = runner1.submit(lambda: {"performance": {"sharpe": 1.2}}, description="hist")
    _wait_done(runner1, task_id)
    runner1.shutdown()

    runner2 = BacktestTaskRunner(max_workers=1)
    state = runner2.peek(task_id)
    assert state is not None
    assert state.status == "done"
    assert state.result["performance"]["sharpe"] == 1.2
    listed = runner2.list_recent(limit=10)
    assert any(s.task_id == task_id for s in listed)
    runner2.shutdown()


def test_runner_persists_failed_task(persisted_env):
    def _boom():
        raise RuntimeError("炸了")

    runner = BacktestTaskRunner(max_workers=1)
    task_id = runner.submit(_boom, description="bad")
    _wait_done(runner, task_id)
    d = ts_mod.get_task_store().load(task_id)
    assert d is not None
    assert d["status"] == "failed"
    assert "RuntimeError" in d["error"]
    runner.shutdown()


# ── REST 导出端点 ──────────────────────────────────────────────────────────────


def _client() -> TestClient:
    from easy_tdx.web import create_app

    app = create_app()
    return TestClient(app)


def test_export_json_and_csv(persisted_env):
    client = _client()
    # 提交一个内联数据回测任务并等待完成
    import numpy as np
    import pandas as pd

    np.random.seed(7)
    n = 200
    close = 10 + np.cumsum(np.random.randn(n) * 0.2 + 0.05)
    dates = pd.date_range("2023-01-01", periods=n, freq="B")
    ohlcv = [
        {
            "datetime": d.strftime("%Y-%m-%d"),
            "open": float(c - 0.05),
            "high": float(c + 0.1),
            "low": float(c - 0.1),
            "close": float(c),
            "vol": 5000.0,
            "amount": float(c * 5000),
        }
        for d, c in zip(dates, close, strict=True)
    ]
    resp = client.post(
        "/api/v1/backtest/run/async",
        json={"strategy": "ma_cross", "params": {"fast": 5, "slow": 20}, "ohlcv": ohlcv},
    )
    assert resp.status_code == 202, resp.text
    task_id = resp.json()["task_id"]
    for _ in range(200):
        st = client.get(f"/api/v1/backtest/tasks/{task_id}").json()
        if st["status"] in ("done", "failed"):
            break
        time.sleep(0.05)
    assert st["status"] == "done", st.get("error")

    # JSON 导出：完整 result
    rj = client.get(f"/api/v1/backtest/tasks/{task_id}/export?format=json")
    assert rj.status_code == 200
    assert "attachment" in rj.headers["content-disposition"]
    assert "performance" in rj.json()

    # CSV 导出：主表（trades 或 performance）
    rc = client.get(f"/api/v1/backtest/tasks/{task_id}/export?format=csv")
    assert rc.status_code == 200
    assert rc.headers["content-type"].startswith("text/csv")
    assert len(rc.text.splitlines()) >= 2


def test_export_rejects_unknown_and_unfinished(persisted_env):
    client = _client()
    assert client.get("/api/v1/backtest/tasks/nope/export").status_code == 400
    resp = client.get("/api/v1/backtest/tasks/nope/export?format=xml")
    # 未知任务先报「未知任务」
    assert resp.status_code == 400


def test_task_list_includes_persisted_history_after_new_app(persisted_env):
    """应用层重启（同 DB）后 /backtest/tasks 仍列出历史任务。"""
    runner = BacktestTaskRunner(max_workers=1)
    task_id = runner.submit(lambda: {"performance": {"total_return": 0.1}}, description="旧任务")
    _wait_done(runner, task_id)
    runner.shutdown()

    client = _client()
    tasks = client.get("/api/v1/backtest/tasks?limit=50").json()["tasks"]
    assert any(t["task_id"] == task_id for t in tasks)


# ── v1.32.6 修复：列表查询懒加载 result_json + 淘汰只移终态 ────────────────────


def test_list_recent_without_results_skips_result_json(persisted_env):
    """include_results=False 时不 SELECT/解析 result_json（列表页瘦身）。

    旧签名无该参数：列表页会把每条任务的完整结果 JSON 拖进内存解析。
    """
    import sqlite3

    from easy_tdx.web.task_store import TaskStore

    store = TaskStore()
    store.save(
        task_id="big",
        status="done",
        created_at=2.0,
        result={
            "performance": {"total_return": 0.25},
            "equity_curve": [{"i": i} for i in range(500)],
        },
    )
    # 把 result_json 打坏：include_results=False 路径根本不读它 → 不受影响
    with sqlite3.connect(store.path) as conn:
        conn.execute("UPDATE backtest_tasks SET result_json = '{not-json' WHERE task_id='big'")
        conn.commit()

    rows = store.list_recent(limit=10, include_results=False)
    assert len(rows) == 1
    assert rows[0]["task_id"] == "big"
    assert rows[0]["status"] == "done"
    assert rows[0]["result"] is None

    # 详情（load）仍取全量并走损坏降级路径
    d = store.load("big")
    assert d is not None and d["result"] is None


def test_eviction_only_removes_terminal_states():
    """超限淘汰只移 done/failed；pending/running 不淘汰（消灭 pending 幽灵）。

    旧实现"淘汰最旧 non-running"会把最早提交、尚未起跑的 pending 条目淘汰，
    其 worker 随后取不到状态直接跳过 → 磁盘遗留 pending 行被恢复成永久
    pending 的幽灵任务。
    """
    from easy_tdx.web.task_runner import BacktestTaskRunner, TaskState

    runner = BacktestTaskRunner(max_workers=1, max_results=2)
    with runner._lock:
        # 插入顺序即 LRU 序：pending 最旧、done 最新
        runner._tasks["p1"] = TaskState(task_id="p1", status="pending")
        runner._tasks["r1"] = TaskState(task_id="r1", status="running")
        runner._tasks["d1"] = TaskState(task_id="d1", status="done")
        runner._evict_if_needed_locked()

    assert "p1" in runner._tasks, "pending 不得被淘汰"
    assert "r1" in runner._tasks, "running 不得被淘汰"
    assert "d1" not in runner._tasks, "超限时应淘汰最旧的终态条目"


def test_eviction_pending_still_evicted_by_old_logic_regression_guard():
    """回归对照：若恢复旧逻辑（淘汰首个 non-running），pending 会先被选中。

    本测试钉死新语义——全部条目为 pending 时宁可不淘汰（超限跳过）。
    """
    from easy_tdx.web.task_runner import BacktestTaskRunner, TaskState

    runner = BacktestTaskRunner(max_workers=1, max_results=1)
    with runner._lock:
        runner._tasks["p1"] = TaskState(task_id="p1", status="pending")
        runner._tasks["p2"] = TaskState(task_id="p2", status="pending")
        runner._evict_if_needed_locked()
    assert len(runner._tasks) == 2  # 无终态可淘汰 → 跳过，不丢任务