dpb/backend/tests/_e2e_weld2.py

"""断链焊接第二段 e2e：配合力纯Python化 + 选择指数(EBV/h²落库+决选) + 亲缘近交。

用法（需在 backend 目录下，PYTHONPATH=.）：
    ENVIRONMENT=dev C:/ai/miniconda3/envs/dpb/python.exe -X utf8 -m tests._e2e_weld2

验证通过后删除脚本产生的数据（ABLUP/配合力/选择指数/决选 记录）。
"""
import asyncio
import os

os.environ.setdefault("ENVIRONMENT", "dev")

from sqlalchemy import delete, select  # noqa: E402

from app.api.v1.module_bre.statistics.service import StatisticsService  # noqa: E402
from app.core.base_schema import AuthSchema, CoreUserSchema  # noqa: E402
from app.core.database import create_async_engine_and_session  # noqa: E402

from app.api.v1.module_system.user.model import UserModel  # noqa: E402  (注册 mapper，避免关系解析失败)
from app.api.v1.module_bre.selection_result.model import SelectionResultModel  # noqa: E402
from app.api.v1.module_bre.statistics.model import (  # noqa: E402
    CombiningAbilityModel,
    PredictionModel,
    PredictionValueModel,
    SelectionIndexModel,
    StatisticsJobModel,
)

AUTH = AuthSchema(user=CoreUserSchema(id=1, is_superuser=True))
TRAIT_ID = 29
TRAIT_CODE = "avg_fruit_weight"
WEIGHTS = {"avg_fruit_weight": 1.0, "max_fruit_weight": 0.8}

created = {"pred": None, "ca": None, "si": None, "job_ids": []}


async def main() -> None:
    _engine, session_factory = create_async_engine_and_session()
    try:
        await _run(session_factory)
        print("E2E 全部通过")
    finally:
        await _engine.dispose()
        await cleanup(session_factory)


async def _run(session_factory) -> None:
    async with session_factory() as db:
        svc = StatisticsService(AUTH, db)

        # 1. ABLUP（含 h²/可靠性，纯 Python 引擎）
        pred_id = await svc.run_ablup(TRAIT_ID, TRAIT_CODE)
        created["pred"] = pred_id
        job = await db.execute(
            select(StatisticsJobModel).where(StatisticsJobModel.result_ref == pred_id, StatisticsJobModel.job_type == "ABLUP")
        )
        created["job_ids"] += [j.id for j in job.scalars().all()]
        preds = await svc.list_predictions()
        p = next(x for x in preds if x.id == pred_id)
        assert p.heritability is not None, "ABLUP 应写入 h²"
        assert 0.0 <= float(p.heritability) <= 1.0, f"h² 越界: {p.heritability}"
        ranking = await svc.ebv_ranking(pred_id)
        assert ranking, "EBV 排行应有数据"
        rels = [float(r.reliability) for r in ranking if r.reliability is not None]
        assert rels and all(0.0 <= r <= 1.0 for r in rels), f"可靠性越界: {rels}"
        print(f"[1] ABLUP ok  pred_id={pred_id} h2={p.heritability} acc={p.accuracy} n={p.train_n} ebv_rows={len(ranking)}")

        # 2. 配合力 GCA/SCA（纯 Python 引擎）
        ca_id = await svc.run_combining(TRAIT_ID, TRAIT_CODE)
        created["ca"] = ca_id
        job = await db.execute(
            select(StatisticsJobModel).where(StatisticsJobModel.result_ref == ca_id, StatisticsJobModel.job_type == "COMBINING")
        )
        created["job_ids"] += [j.id for j in job.scalars().all()]
        ca = await svc.combining_detail(ca_id)
        assert ca.gca_json and ca.sca_json, "GCA/SCA 应有结果"
        print(f"[2] COMBINING ok  ca_id={ca_id} parents={len(ca.gca_json)} combos={len(ca.sca_json)}")

        # 3. 选择指数（EBV 接入 + h² 加权 + 落库）
        si = await svc.selection_index(weights=WEIGHTS, batch_ids={TRAIT_CODE: pred_id}, use_h2=True, top_n=5)
        created["si"] = si["id"]
        assert si["top"], "选择指数应有排名"
        assert si["heritabilities"][TRAIT_CODE] is not None, "应带入批次 h²"
        top = await svc.list_index_batches()
        assert any(x.id == si["id"] for x in top), "选择指数批次应落库可查"
        print(f"[3] SELECTION-INDEX ok  si_id={si['id']} method={si['method']} h2={si['heritabilities']} top={len(si['top'])}")

        # 4. 写入决选
        applied = await svc.apply_selection(si["id"], top_n=3)
        assert applied["written"] >= 1, "决选应写入"
        print(f"[4] APPLY ok  written={applied['written']} year={applied['selection_year']}")

        # 5. 亲缘/近交
        kin = await svc.kinship_matrix(threshold=0.25)
        assert kin["n"] > 0 and kin["individuals"], "亲缘分析应有结果"
        print(f"[5] KINSHIP ok  n={kin['n']} flags={len(kin['flags'])} F>0={len(kin['inbreeding'])} matrix={'yes' if kin['matrix'] else 'no'}")
        if kin["flags"]:
            print("    top flags:", kin["flags"][:3])


async def cleanup(session_factory) -> None:
    async with session_factory() as db:
        si_id = created["si"]
        if si_id:
            await db.execute(
                delete(SelectionResultModel).where(SelectionResultModel.reason.like(f"选择指数批次#{si_id}%"))
            )
            await db.execute(delete(SelectionIndexModel).where(SelectionIndexModel.id == si_id))
        if created["ca"]:
            await db.execute(delete(CombiningAbilityModel).where(CombiningAbilityModel.id == created["ca"]))
        if created["pred"]:
            await db.execute(delete(PredictionValueModel).where(PredictionValueModel.prediction_id == created["pred"]))
            await db.execute(delete(PredictionModel).where(PredictionModel.id == created["pred"]))
        if created["job_ids"]:
            await db.execute(delete(StatisticsJobModel).where(StatisticsJobModel.id.in_(created["job_ids"])))
        await db.commit()
        print("已清理 e2e 产生的数据")


if __name__ == "__main__":
    asyncio.run(main())