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。
"""
import datetime as dt
import sys
import traceback

import pool

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"]


# ---------------------------------------------------------------- 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()