akg-factor-bridge/test_plan_snapshot.py

"""接口改读当日全量快照（2026-09-10，台账 057）：截取规则、名次重编、回落、字段一致。

要解决的问题：计划接口每次请求都现场装配整池，实测三十三秒。而计划是日频的、用的全是
昨收与昨夜数据，出计划时本来就落了一份全量快照（不裁剪、不设主题限额），接口一直没读它。

离线，不连库不起服务。跑法：python3 test_plan_snapshot.py，预期最后一行是 ALL OK。
真实快照存在时会额外做一轮逐行比对，不存在就跳过那一轮并说明。
"""
import glob
import json
import os
import sys
import types

for _n in ("pandas", "pymysql", "psycopg"):
    if _n not in sys.modules:
        try:
            __import__(_n)
        except Exception:  # noqa: BLE001
            _m = types.ModuleType(_n)
            if _n == "pandas":
                _m.DataFrame = type("DataFrame", (), {})
            sys.modules[_n] = _m

import config      # noqa: E402
import plan        # noqa: E402

FAILED = []


def t(name, cond, extra=""):
    print(("  ok  " if cond else "  FAIL ") + name + (("   " + str(extra)) if not cond else ""))
    if not cond:
        FAILED.append(name)


def _row(code, theme, rank):
    return {"code": code, "rank": rank, "score": 100.0 - rank,
            "evidence": ({"theme": theme} if theme else None)}


def test_pick_rules():
    print("[截取规则]")
    rows = [_row("A", "光伏", 1), _row("B", "光伏", 2), _row("C", "光伏", 3),
            _row("D", "军工", 4), _row("E", None, 5)]
    t("不设限就按顺序取", [r["code"] for r in plan.pick_rows(rows, 3, 0)] == ["A", "B", "C"])
    got = plan.pick_rows(rows, 4, 2)
    t("每个主题最多两条，第三条光伏被跳过",
      [r["code"] for r in got] == ["A", "B", "D", "E"], [r["code"] for r in got])
    t("没有主题的行归到同一个桶，不当成各自独立",
      [r["code"] for r in plan.pick_rows(
          [_row("X", None, 1), _row("Y", None, 2), _row("Z", "光伏", 3)], 3, 1)] == ["X", "Z"])
    t("要的条数比总行数多就全给", len(plan.pick_rows(rows, 99, 0)) == 5)
    t("要零条时先给一条（照抄原实现的边界：先放进去再判断够没够，不是我们要它这样）",
      len(plan.pick_rows(rows, 0, 0)) == 1)
    t("要一条就给一条", [r["code"] for r in plan.pick_rows(rows, 1, 0)] == ["A"])


def test_replay_against_real_snapshot():
    print("[拿真实快照逐行比对 —— 这条是这次改动成立的根据]")
    fs = sorted(glob.glob(os.path.join(config.PLAN_SNAPSHOT_DIR, "plan_*.json")))
    if not fs:
        print("  --  本机没有快照文件，跳过这一轮（在部署机上跑才有）")
        return
    p = fs[-1]
    snap = json.load(open(p, encoding="utf-8"))
    sp = snap.get("shown_params") or {}
    if not sp or not isinstance(snap.get("main"), list):
        print("  --  最新快照没有全量段或落盘参数，跳过：" + os.path.basename(p))
        return
    print("  用 %s（落盘参数 %s）" % (os.path.basename(p), sp))
    for seg, shown_key, n_key in (("main", "main_shown", "top"),
                                  ("observe", "observe_shown", "obs_top")):
        got = plan.pick_rows(snap[seg], sp[n_key], sp["theme_cap"])
        shown = snap[shown_key]
        t("%s 截出来的代码序列与当时下发的完全一致" % seg,
          [r["code"] for r in got] == [r["code"] for r in shown],
          ([r["code"] for r in got][:6], [r["code"] for r in shown][:6]))
        # 名次是唯一该不一样的字段：全量快照里存的是全量排名，下发时要重编成连续名次
        diff = {k for a, b in zip(got, shown) for k in set(a) | set(b) if a.get(k) != b.get(k)}
        t("%s 除名次外逐字段一致" % seg, diff <= {"rank"}, diff)
        renum = [{**r, "rank": i} for i, r in enumerate(got, 1)]
        t("%s 名次重编后与当时下发的一模一样" % seg, renum == shown)


def test_from_snapshot(tmpdir=None):
    print("[from_snapshot 的输出与回落]")
    fs = sorted(glob.glob(os.path.join(config.PLAN_SNAPSHOT_DIR, "plan_*.json")))
    if not fs:
        print("  --  本机没有快照文件，跳过这一轮")
        return
    ds = os.path.basename(fs[-1])[len("plan_"):-len(".json")]
    d = plan.from_snapshot(ds, top=5, obs_top=3, theme_cap=2)
    t("名次是连续的 1 到 n", [r["rank"] for r in d["main"]] == list(range(1, len(d["main"]) + 1)))
    t("条数不超过请求的", len(d["main"]) <= 5 and len(d["observe"]) <= 3)
    t("主题限额按本次请求生效，不是落盘那次的", d["theme_cap"] == 2)
    t("标明了这份是从快照来的", d.get("plan_source") == "snapshot")
    t("带上快照的生成时刻，一眼看得出多新", bool(d.get("snapshot_generated_at")))
    for k in ("main_shown", "observe_shown", "shown_params"):
        t("不把落盘那次的裁剪结果带出去（%s）" % k, k not in d)
    for k in ("market", "regime"):
        t("不把快照里的 %s 带出去（接口每次实时读）" % k, k not in d)
    for k in ("date", "generated_at", "plan_version", "counts", "candidates", "watch",
              "segments_pointed", "card_counts", "card_params", "changes", "encoding"):
        t("顶层保留 %s" % k, k in d)
    # 全量计数不随裁剪变化：counts 说的是整池有多少，不是这次给了几条
    t("全量计数原样保留", (d.get("counts") or {}).get("main", 0) >= len(d["main"]))
    # 认不出的日期、对不上的日期一律抛异常，交给调用方回落
    try:
        plan.from_snapshot("1999-01-01", 5, 3, 2)
        t("快照不存在要抛异常", False)
    except Exception as e:  # noqa: BLE001
        t("快照不存在要抛异常", isinstance(e, (FileNotFoundError, RuntimeError)), type(e).__name__)


def test_api_falls_back():
    print("[接口回落：读不到快照不能让接口挂掉]")
    src = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "api.py"),
               encoding="utf-8").read()
    t("有总开关，能一键回到每次现算的老路", "PLAN_FROM_SNAPSHOT" in src)
    t("快照那条路包了兜底，任何异常都回落", "except Exception" in src and "回落到实时装配" in src)
    t("回落一定记一行警告说明原因（不然只会表现成偶尔变慢）",
      "_access.warning" in src and "快照读取失败" in src)
    i_try = src.index("plan.from_snapshot")
    i_live = src.index("plan.collect(date, top, obs_top, theme_cap)")
    t("先试快照、后现算", i_try < i_live)
    t("访问日志写明这次走的哪条路", "source=%s" in src)


def main():
    test_pick_rules()
    test_replay_against_real_snapshot()
    test_from_snapshot()
    test_api_falls_back()
    print("-" * 70)
    if FAILED:
        print("FAILED %d: %s" % (len(FAILED), "; ".join(FAILED)))
        return 1
    print("ALL OK — 计划读快照：截取规则 / 与真实快照逐行一致 / 输出与回落 / 接口接线")
    return 0


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