akg-factor-bridge/test_plan_cache.py

"""选股计划的接口缓存（2026-09-10 故障修复）：数据日换了立刻失效、同键只算一次、强刷绕开缓存。

09-10 早上页面报「上游的选股计划读不到 · timeout of 60000ms exceeded」，根因两半：
缓存只活 10 分钟，盘前 08:40 预热那份 08:50 就过期；而装配又在锁外面，09:02 页面几路
并发进来各跑一遍三十多秒的装配、互相抢库连接，叠加破了 60 秒。这里把两条都钉住。

离线，不连库不起服务。跑法：python3 test_plan_cache.py，预期最后一行是 ALL OK。
"""
import sys
import threading
import time
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 api            # noqa: E402
import plan           # noqa: E402


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


class _Spy:
    """替身：记录 collect 被调了几次、每次算多久。"""

    def __init__(self, delay=0.0):
        self.calls = []
        self.delay = delay
        self._lock = threading.Lock()

    def collect(self, date, top, obs_top, theme_cap):
        with self._lock:
            self.calls.append((date, top, obs_top, theme_cap))
        time.sleep(self.delay)
        return {"date": date or "2026-09-09", "main": [], "observe": [], "_full": "该被丢掉"}


def _reset():
    api._PLAN_CACHE.clear()
    api._PLAN_BUILD_LOCKS.clear()


def test_hit_and_miss():
    print("[命中与未命中]")
    spy = _Spy()
    _reset()
    old_collect, old_latest = plan.collect, plan._latest_date
    plan.collect = spy.collect
    plan._latest_date = lambda tbl: "2026-09-09"
    try:
        d1, hit1 = api._plan_cached(None, 300, 100, 2, False)
        d2, hit2 = api._plan_cached(None, 300, 100, 2, False)
        t("第一次未命中、第二次命中", hit1 is False and hit2 is True)
        t("只算了一次", len(spy.calls) == 1, spy.calls)
        t("内部大字段被丢掉，不进缓存也不出接口", "_full" not in d1 and "_full" not in d2)

        # 参数不同 = 不同的键。PMS 用 300/100，人工联调常用 30/20，两份互不干扰
        api._plan_cached(None, 30, 20, 2, False)
        t("参数不同要各算各的", len(spy.calls) == 2, spy.calls)

        # 数据日换了，旧那份必须立刻失效，不能等缓存到期
        plan._latest_date = lambda tbl: "2026-09-10"
        _d, hit3 = api._plan_cached(None, 300, 100, 2, False)
        t("上游换日后旧缓存立刻失效", hit3 is False and len(spy.calls) == 3)

        # 强刷绕开缓存，但要把新结果写回去
        n = len(spy.calls)
        _d, hit4 = api._plan_cached(None, 300, 100, 2, True)
        t("强刷绕开缓存", hit4 is False and len(spy.calls) == n + 1)
        _d, hit5 = api._plan_cached(None, 300, 100, 2, False)
        t("强刷之后缓存是新的", hit5 is True and len(spy.calls) == n + 1)
    finally:
        plan.collect, plan._latest_date = old_collect, old_latest


def test_single_flight():
    print("[同键只算一次 —— 09-10 超时的那一半]")
    spy = _Spy(delay=0.4)          # 装配很慢，真实环境是三十多秒
    _reset()
    old_collect, old_latest = plan.collect, plan._latest_date
    plan.collect = spy.collect
    plan._latest_date = lambda tbl: "2026-09-09"
    try:
        results, errs = [], []

        def one():
            try:
                results.append(api._plan_cached(None, 300, 100, 2, False))
            except Exception as e:  # noqa: BLE001
                errs.append(e)

        ths = [threading.Thread(target=one) for _ in range(8)]
        t0 = time.time()
        for x in ths:
            x.start()
        for x in ths:
            x.join()
        el = time.time() - t0

        t("八路并发一个都没出错", not errs, errs)
        t("八路都拿到了结果", len(results) == 8, len(results))
        t("装配只跑了一遍（原来是八遍）", len(spy.calls) == 1, spy.calls)
        t("总耗时接近单跑一次，不是八次叠加", el < 0.4 * 3, "%.2f 秒" % el)
        t("只有一路是未命中，其余都吃了缓存",
          sum(1 for _d, h in results if not h) == 1, [h for _d, h in results])
    finally:
        plan.collect, plan._latest_date = old_collect, old_latest


def test_different_keys_not_blocked():
    print("[不同参数互不阻塞]")
    spy = _Spy(delay=0.4)
    _reset()
    old_collect, old_latest = plan.collect, plan._latest_date
    plan.collect = spy.collect
    plan._latest_date = lambda tbl: "2026-09-09"
    try:
        def one(top):
            api._plan_cached(None, top, 100, 2, False)
        ths = [threading.Thread(target=one, args=(x,)) for x in (300, 30, 50)]
        t0 = time.time()
        for x in ths:
            x.start()
        for x in ths:
            x.join()
        el = time.time() - t0
        t("三种参数各算一次", len(spy.calls) == 3, spy.calls)
        t("是并行不是排队", el < 0.4 * 2.5, "%.2f 秒" % el)
    finally:
        plan.collect, plan._latest_date = old_collect, old_latest


def test_ttl_covers_trading_day():
    print("[存活时间要盖住整个交易日]")
    t("默认至少 8 小时（08:40 预热要能管到收盘）",
      api.PLAN_CACHE_SEC >= 8 * 3600, api.PLAN_CACHE_SEC)
    t("锁字典与缓存字典一起清理，不会只涨不落",
      "_PLAN_BUILD_LOCKS.pop" in open(api.__file__.replace(".pyc", ".py"), encoding="utf-8").read())


def main():
    test_hit_and_miss()
    test_single_flight()
    test_different_keys_not_blocked()
    test_ttl_covers_trading_day()
    print("ALL OK — 选股计划缓存：命中与失效 / 同键只算一次 / 不同参数不互相阻塞 / 存活盖住交易日")


if __name__ == "__main__":
    main()