akg-factor-bridge/test_pool_logic.py

# -*- coding: utf-8 -*-
"""入池/留池/出池纯逻辑的单测（不连任何库）。

运行:  docker compose exec -T akg-factor-bridge python test_pool_logic.py
全过输出 "ALL PASS (n cases)"，任一失败退出码 1。
被测函数: pool.decide / pool.build_remark / pool.build_recycle_docs /
pool.segment_started_ratios / pool.strategy_context_entry（后两个 2026-09-03 加）。

开发机没有 pandas 与数据库驱动时也能跑：只给缺席的模块装最小桩（与 test_market_context.py
同一约定，模块存在时不覆盖），被测函数本身不碰这些依赖。
"""
import datetime as dt
import sys
import traceback
import types

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

import pool  # noqa: E402

RESULTS = []


def case(name):
    def deco(fn):
        RESULTS.append((name, fn))
        return fn
    return deco


TODAY = "2026-08-04"


@case("首次建池: 空池 + 计划20只 + 持仓3只(1只重叠) → 池=计划∪持仓, 无出池")
def _():
    plan = [f"SH60{i:04d}" for i in range(20)]
    hold = {"SH600000", "SZ000001", "SZ000002"}          # SH600000 与计划重叠
    d = pool.decide(set(), {}, plan, hold, set(), 60, TODAY)
    assert len(d["pool"]) == 22, d["pool"]
    assert set(plan) <= set(d["pool"]) and hold <= set(d["pool"])
    assert d["retained_holdings"] == ["SZ000001", "SZ000002"]
    assert d["recycled"] == [] and d["cap_evicted"] == [] and d["observers"] == []
    assert d["new_entrants"] == sorted(plan)
    # 成员记录: 计划票有 last_plan, 纯持仓票只有 added
    assert d["meta"]["SH600000"]["last_plan"] == TODAY
    assert "last_plan" not in d["meta"]["SZ000001"]


@case("掉榜不等于变坏: 掉出计划且无持仓、形态未恶化 → 留池观察")
def _():
    old = {"SH600001", "SH600002", "SH600003"}
    meta = {c: {"added": "2026-08-01", "last_plan": "2026-08-01"} for c in old}
    d = pool.decide(old, meta, ["SH600001"], set(), set(), 60, TODAY)
    assert d["observers"] == ["SH600002", "SH600003"], d
    assert set(d["pool"]) == old                      # 都还在池里
    assert d["recycled"] == []


@case("出池三条件缺一不可: 无持仓 + 不在计划 + 形态恶化 → 回收站")
def _():
    old = {"SH600001", "SH600002", "SH600003", "SH600004"}
    meta = {c: {"added": "2026-08-01"} for c in old}
    bad = {"SH600002", "SH600003", "SH600004"}
    #  600002 恶化但仍在计划 → 留; 600003 恶化但有持仓 → 留(警示); 600004 三条全中 → 回收
    d = pool.decide(old, meta, ["SH600001", "SH600002"], {"SH600003"}, bad, 60, TODAY)
    assert d["recycled"] == ["SH600004"], d
    assert "SH600002" in d["pool"] and "SH600003" in d["pool"]
    assert d["held_bad"] == ["SH600003"]


@case("池满出清: 只清留池观察、按最久没上榜的先清, 计划与持仓永不清")
def _():
    plan = [f"SH61{i:04d}" for i in range(5)]
    hold = {"SZ000001"}
    old = set(plan) | hold | {"SH620001", "SH620002", "SH620003"}
    meta = {"SH620001": {"added": "2026-07-01", "last_plan": "2026-07-10"},
            "SH620002": {"added": "2026-07-01", "last_plan": "2026-07-20"},
            "SH620003": {"added": "2026-07-01", "last_plan": "2026-07-30"}}
    d = pool.decide(old, meta, plan, hold, set(), 8, TODAY)     # 5+1+3=9 > 8, 清 1 只
    assert d["cap_evicted"] == ["SH620001"], d                  # 最久没上榜的
    assert len(d["pool"]) == 8
    assert d["recycled"] == []                                  # 池满出清不进回收站
    assert "SH620001" not in d["meta"]                          # 出池即清成员记录


@case("计划+持仓本身超上限时不硬砍 (允许超, 只把观察清空)")
def _():
    plan = [f"SH63{i:04d}" for i in range(10)]
    hold = {f"SZ00{i:04d}" for i in range(5)}
    old = set(plan) | hold | {"SH640001"}
    d = pool.decide(old, {"SH640001": {"added": "2026-07-01"}}, plan, hold, set(), 12, TODAY)
    assert d["cap_evicted"] == ["SH640001"]
    assert len(d["pool"]) == 15                                 # 10+5, 超 12 但不砍计划/持仓


@case("恶化名单为空 (结论读不到的降级) → 一只都不回收")
def _():
    old = {"SH600001", "SH600002"}
    d = pool.decide(old, {}, [], set(), set(), 60, TODAY)
    assert d["recycled"] == [] and set(d["observers"]) == old


@case("幂等: 同一天跑两遍, 第二遍无新进无回收, 池子不变")
def _():
    plan = ["SH600001", "SH600002"]
    d1 = pool.decide(set(), {}, plan, {"SZ000001"}, set(), 60, TODAY)
    d2 = pool.decide(set(d1["pool"]), d1["meta"], plan, {"SZ000001"}, set(), 60, TODAY)
    assert d2["pool"] == d1["pool"]
    assert d2["new_entrants"] == [] and d2["recycled"] == []


@case("remark 汇总句: 数字对得上, 持仓恶化有警示, 降级有说明")
def _():
    d = pool.decide({"SH600009"}, {"SH600009": {"added": "2026-08-01"}},
                    ["SH600001"], {"SH600003"}, {"SH600003"}, 60, TODAY)
    r = pool.build_remark(d, "2026-08-04", "2026-08-04 07:20:00", degraded="测试降级说明")
    # 池 = SH600001(计划) + SH600003(持仓, 恶化但保留) + SH600009(未恶化留池) = 3 只
    assert "共入池 3 只" in r["summary"], r["summary"]
    assert "当日计划 1 只" in r["summary"] and "持仓保留 1 只" in r["summary"]
    assert "留池观察 1 只" in r["summary"], r["summary"]
    assert "警示" in r["summary"] and "SH600003" in r["summary"]
    assert "测试降级说明" in r["summary"]
    assert r["retained_positions"] == ["SH600003"]
    assert r["factor_details"][0]["selected_codes"] == ["SH600001"]
    assert r["factor_details"][0]["trade_date"] == "2026-08-04"


@case("回收站文档: 字段与现有格式一致 (group_id/removal_batch/removed_at/stock_code)")
def _():
    d = pool.decide({"SH600004"}, {}, [], set(), {"SH600004"}, 60, TODAY)
    now = dt.datetime(2026, 8, 4, 7, 20, 0)
    docs = pool.build_recycle_docs(d, "6865e8eb97623ef14325c8d3", "AKG每日选股计划池",
                                   "489281497140", now)
    assert len(docs) == 1
    doc = docs[0]
    for k in ("group_id", "group_name", "org_id", "removal_batch", "removed_at",
              "stock_code", "reason"):
        assert k in doc, k
    assert doc["stock_code"] == "SH600004"
    assert doc["removal_batch"] == now.isoformat()
    assert doc["removed_at"] == now
    assert "形态恶化" in doc["reason"]


@case("入池上下文: 2026-09-03 补的证据字段只加不改, 环节已启动比例优先用数据基座台账的 moved")
def _():
    segs = [{"segment": "散热器件", "members_total": 20, "moved": 5, "started_count": 4},
            {"segment": "液冷", "members_total": 10, "moved": None, "started_count": 3},
            {"segment": "无成员数", "members_total": None, "moved": 2, "started_count": 2}]
    ratio = pool.segment_started_ratios(segs)
    assert ratio == {"散热器件": 0.25, "液冷": 0.3, "无成员数": None}, ratio
    row = {"code": "SH600000", "score": 231.5, "tier": "强传导", "upside": 0.25,
           "evidence": {"theme": "散热器件", "n_sources": 2, "moved_ratio": 0.25},
           "verdict": "候选", "reasons": ["a", "b", "c", "d", "e"], "card_rank": 1,
           "card": {"pct0": 4.1, "chain_fit": 0.8, "accum": {"state": "明确吸筹·量在价先", "age": 3}}}
    e = pool.strategy_context_entry(row, "2026-09-03", ratio)
    for k, v in {"factor_code": "akg_score", "score": 231.5, "tier": "强传导", "upside": 0.25,
                 "theme": "散热器件", "plan_date": "2026-09-03", "verdict": "候选", "card_rank": 1,
                 "chain_fit": 0.8, "n_sources": 2, "pct0": 4.1, "accum_state": "明确吸筹·量在价先",
                 "accum_age": 3, "segment_started_ratio": 0.25}.items():
        assert e[k] == v, (k, e[k])
    assert e["reasons"] == ["a", "b", "c", "d"]                 # 理由仍只留四条
    assert "source" not in e and "plan_version" not in e         # 这两个键由调用方补
    # 证据线缺席的行 (无 card / 无 evidence): 新字段为 None, 不抛错
    e = pool.strategy_context_entry({"code": "SZ000001", "score": 100.0}, "2026-09-03", ratio)
    assert e["theme"] is None and e["chain_fit"] is None and e["segment_started_ratio"] is None
    assert e["accum_state"] is None and e["pct0"] is None and e["reasons"] == []


# ---------------------------------------------------------------- runner
def main():
    passed, failed = 0, 0
    for name, fn in RESULTS:
        try:
            fn()
            print(f"  PASS  {name}")
            passed += 1
        except Exception:
            print(f"  FAIL  {name}")
            traceback.print_exc()
            failed += 1
    print("-" * 60)
    if failed:
        print(f"FAILED: {failed} / {passed + failed}")
        sys.exit(1)
    print(f"ALL PASS ({passed} cases)")


if __name__ == "__main__":
    main()