dpb/backend/scripts/test_bre_analysis_mendelian_tc.py

# -*- coding: utf-8 -*-
"""孟德尔/亲本一致性检验正式 tc 套件：TestClient 走真实 API + 落库/隔离 golden 断言。

fixture（镜像 e2e_mendelian_persist）：P1/P2 全 0/1 → T1、P3/P4 全 0/0 → T2（5 标记）。
T1 全 0/0 → rate=1.0 兼容；T2 全 1/1 → 每标记 alien → rate=0.0 → flag。
经真实 HTTP 端点 POST /api/v1/bre/genotype_call/mendelian/check/run 断言：
  [1] 返回 check_id / isolated=1 / cleared=0 + 计数 2/1/0/0
  [2] bre_mendelian_check 批次落库（阈值/计数/data_version/input_hash）
  [3] 子表 2 行：T1 rate=1.0 flag=false / T2 rate=0.0 flag=true incompatible=5
  [4] T2 进入隔离注册表 bre_pedigree_exclusion（source=mendelian / source_check_id）
  [5] 阈值 0.0 重跑 → flagged=0 / cleared=1 + 隔离行清除（可逆）
  [6] GET 只读预览零回归（无 check_id / 纯计算）
依赖: Redis + PG 正常（TestClient 走真实 lifespan）。运行后自动清理。
"""
import os
os.environ["ENVIRONMENT"] = "dev"
os.environ["PYTHONUTF8"] = "1"

import sys, asyncio  # noqa: E402
sys.path.insert(0, r"d:\dpb\dpb\backend")

import main  # noqa: E402
from fastapi.testclient import TestClient  # noqa: E402
from sqlalchemy import delete, select  # 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.target.model import TargetModel  # noqa: E402
from app.api.v1.module_bre.trait.model import TraitModel  # noqa: E402
from app.api.v1.module_bre.germplasm.model import BreedingGermplasmModel  # noqa: E402
from app.api.v1.module_bre.cross_combination.model import CrossCombinationModel  # noqa: E402
from app.api.v1.module_bre.tree.model import TreeModel, TreePedigreeExclusionModel  # noqa: E402
from app.api.v1.module_bre.trait_observation.model import TraitObservationModel  # noqa: E402
from app.api.v1.module_bre.marker.model import MarkerModel  # noqa: E402
from app.api.v1.module_bre.genotype_dataset.model import GenotypingDatasetModel  # noqa: E402
from app.api.v1.module_bre.genotype_sample.model import GenotypeSampleModel  # noqa: E402
from app.api.v1.module_bre.genotype_call.model import (  # noqa: E402
    GenotypeCallModel,
    MendelianCheckModel,
    MendelianCheckTreeModel,
)
from app.api.v1.module_bre.genotype_call.service import GENOTYPE_DATA_VERSION  # noqa: E402

create_app = main.create_app
TOKEN = None
ok, fail = 0, 0

PREFIX = "TC_MEND"
tokens: dict[str, list[int]] = {
    "germ": [], "tree": [], "combo": [], "target": [],
    "marker": [], "dataset": [], "sample": [], "call": [],
    "obs": [],
}
check_ids: list[int] = []
FIX: dict = {}


def check(name, cond, detail=""):
    global ok, fail
    if cond:
        ok += 1
        print(f"  [ok] {name} {detail}")
    else:
        fail += 1
        print(f"  [FAIL] {name} {detail}")


def login(client):
    global TOKEN
    d = {"username": "super", "password": "123456", "grant_type": "password", "login_type": "PC端"}
    r = client.post("/api/v1/system/auth/login", data=d)
    b = r.json()
    if r.status_code == 200 and b.get("code") == 0:
        TOKEN = b["data"]["access_token"]
        return
    key = client.get("/api/v1/system/auth/captcha/get").json()["data"]["key"]
    client.post("/api/v1/system/auth/captcha/slider/complete", json={"captcha_key": key})
    d["captcha_key"] = key
    r = client.post("/api/v1/system/auth/login", data=d)
    b = r.json()
    assert r.status_code == 200 and b.get("code") == 0, f"LOGIN FAIL {r.status_code} {b}"
    TOKEN = b["data"]["access_token"]


def auth():
    return {"Authorization": f"Bearer {TOKEN}"}


async def _build_fixture() -> None:
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            target = TargetModel(target_name=f"{PREFIX}-T", created_id=1)
            db.add(target)
            await db.flush()
            tokens["target"].append(target.id)
            trait = TraitModel(trait_code=f"tM_{PREFIX}", trait_name=f"果径{PREFIX}",
                               data_type="numeric", unit="mm", is_core="1", direction="desc",
                               into_ebv="1", default_h2=0.5, created_id=1)
            db.add(trait)
            await db.flush()
            P1 = BreedingGermplasmModel(cultivar_name=f"{PREFIX}-P1", can_be_female=True, created_id=1)
            P2 = BreedingGermplasmModel(cultivar_name=f"{PREFIX}-P2", can_be_male=True, created_id=1)
            P3 = BreedingGermplasmModel(cultivar_name=f"{PREFIX}-P3", can_be_female=True, created_id=1)
            P4 = BreedingGermplasmModel(cultivar_name=f"{PREFIX}-P4", can_be_male=True, created_id=1)
            db.add_all([P1, P2, P3, P4])
            await db.flush()
            tokens["germ"] += [P1.id, P2.id, P3.id, P4.id]
            C0 = CrossCombinationModel(combination_code=f"{PREFIX}-C0", bre_target_id=target.id,
                                       female_parent_id=P1.id, male_parent_id=P2.id,
                                       design_type="full_diallel", created_id=1)
            C1 = CrossCombinationModel(combination_code=f"{PREFIX}-C1", bre_target_id=target.id,
                                       female_parent_id=P3.id, male_parent_id=P4.id,
                                       design_type="full_diallel", created_id=1)
            db.add_all([C0, C1])
            await db.flush()
            tokens["combo"] += [C0.id, C1.id]
            T1 = TreeModel(combination_id=C0.id, tree_no=f"{PREFIX}-T1", status="alive",
                           stage="evaluation", generation="F1", dam_id=P1.id, sire_id=P2.id, created_id=1)
            T2 = TreeModel(combination_id=C1.id, tree_no=f"{PREFIX}-T2", status="alive",
                           stage="evaluation", generation="F1", dam_id=P3.id, sire_id=P4.id, created_id=1)
            db.add_all([T1, T2])
            await db.flush()
            tokens["tree"] += [T1.id, T2.id]
            FIX.update({"T1": T1.id, "T2": T2.id, "P1": P1.id, "P2": P2.id,
                        "P3": P3.id, "P4": P4.id})
            o1 = TraitObservationModel(tree_id=T1.id, combination_id=C0.id, trait_id=trait.id,
                                       value_numeric=25.0, evaluate_year=2025, created_id=1)
            o2 = TraitObservationModel(tree_id=T2.id, combination_id=C1.id, trait_id=trait.id,
                                       value_numeric=20.0, evaluate_year=2025, created_id=1)
            db.add_all([o1, o2])
            await db.flush()
            tokens["obs"] += [o1.id, o2.id]
            markers = []
            for j in range(5):
                m = MarkerModel(marker_name=f"{PREFIX}-M{j + 1}", marker_type="SNP",
                                chromosome=str(j), position=j * 100, created_id=1)
                db.add(m)
                markers.append(m)
            await db.flush()
            tokens["marker"] = [m.id for m in markers]
            mids = [m.id for m in markers]
            ds = GenotypingDatasetModel(dataset_name=f"{PREFIX}-DS", platform="SNP",
                                        purpose="MD", created_id=1)
            db.add(ds)
            await db.flush()
            tokens["dataset"].append(ds.id)
            FIX["dataset_id"] = ds.id

            def mk_sample(ds_id, no, source_type, source_id):
                s = GenotypeSampleModel(sample_name=no, dataset_id=ds_id,
                                        source_type=source_type, source_id=source_id, created_id=1)
                db.add(s)
                return s

            samp = [
                mk_sample(ds.id, f"{PREFIX}-P1-S", "germplasm", P1.id),
                mk_sample(ds.id, f"{PREFIX}-P2-S", "germplasm", P2.id),
                mk_sample(ds.id, f"{PREFIX}-P3-S", "germplasm", P3.id),
                mk_sample(ds.id, f"{PREFIX}-P4-S", "germplasm", P4.id),
                mk_sample(ds.id, f"{PREFIX}-T1-S", "tree", T1.id),
                mk_sample(ds.id, f"{PREFIX}-T2-S", "tree", T2.id),
            ]
            await db.flush()
            tokens["sample"] = [s.id for s in samp]
            calls = []
            gt_map = {"P1": "0/1", "P2": "0/1", "P3": "0/0", "P4": "0/0", "T1": "0/0", "T2": "1/1"}
            for name, s in zip(("P1", "P2", "P3", "P4", "T1", "T2"), samp):
                for mid in mids:
                    calls.append(GenotypeCallModel(sample_id=s.id, marker_id=mid,
                                                   allele=gt_map[name], created_id=1))
            db.add_all(calls)
            await db.flush()
            tokens["call"] = [c.id for c in calls]
            await db.commit()
    finally:
        await engine.dispose()


async def _verify_and_cleanup() -> None:
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            r = FIX["run1"]
            # ---- [2] 批次落库 ----
            crow = await db.get(MendelianCheckModel, r["check_id"])
            check("[2] bre_mendelian_check 落库（阈值/计数/data_version/input_hash）",
                  crow is not None and float(crow.compat_threshold) == 0.98
                  and crow.trees_checked == 2 and crow.trees_flagged == 1
                  and crow.data_version == GENOTYPE_DATA_VERSION and crow.input_hash,
                  f"v={crow.data_version if crow else None}")
            # ---- [3] 子表 ----
            childs = (await db.execute(select(MendelianCheckTreeModel).where(
                MendelianCheckTreeModel.check_id == r["check_id"]))).scalars().all()
            child_by_tree = {int(c.tree_id): c for c in childs}
            check("[3] 子表 2 行：T1 rate=1.0 flag=false / T2 rate=0.0 flag=true incompatible=5",
                  len(childs) == 2
                  and float(child_by_tree[FIX["T1"]].compatibility_rate) == 1.0
                  and not child_by_tree[FIX["T1"]].flag
                  and float(child_by_tree[FIX["T2"]].compatibility_rate) == 0.0
                  and child_by_tree[FIX["T2"]].flag
                  and child_by_tree[FIX["T2"]].status == "checked"
                  and child_by_tree[FIX["T2"]].incompatible == 5,
                  f"T1={child_by_tree[FIX['T1']].compatibility_rate}/{child_by_tree[FIX['T1']].flag} "
                  f"T2={child_by_tree[FIX['T2']].compatibility_rate}/{child_by_tree[FIX['T2']].flag}")
            # ---- [4] 隔离注册表在 main_ 中 run1 后立即验（run2 会清除），此处只验 [5] ----
            # ---- [5] 可逆：隔离行清除 ----
            if FIX.get("run2"):
                excl2 = (await db.execute(select(TreePedigreeExclusionModel).where(
                    TreePedigreeExclusionModel.is_deleted.is_(False),
                    TreePedigreeExclusionModel.tree_id == FIX["T2"]))).scalars().all()
                check("[5] 阈值 0.0 后 T2 隔离行已清除", len(excl2) == 0, f"n={len(excl2)}")

            # ---- cleanup ----
            if tokens["tree"]:
                await db.execute(delete(TreePedigreeExclusionModel).where(
                    TreePedigreeExclusionModel.tree_id.in_(tokens["tree"])))
            if check_ids:
                await db.execute(delete(MendelianCheckTreeModel).where(
                    MendelianCheckTreeModel.check_id.in_(check_ids)))
                await db.execute(delete(MendelianCheckModel).where(
                    MendelianCheckModel.id.in_(check_ids)))
            if tokens["call"]:
                await db.execute(delete(GenotypeCallModel).where(
                    GenotypeCallModel.id.in_(tokens["call"])))
            if tokens["sample"]:
                await db.execute(delete(GenotypeSampleModel).where(
                    GenotypeSampleModel.id.in_(tokens["sample"])))
            if tokens["marker"]:
                await db.execute(delete(MarkerModel).where(MarkerModel.id.in_(tokens["marker"])))
            if tokens["dataset"]:
                await db.execute(delete(GenotypingDatasetModel).where(
                    GenotypingDatasetModel.id.in_(tokens["dataset"])))
            if tokens["obs"]:
                await db.execute(delete(TraitObservationModel).where(
                    TraitObservationModel.id.in_(tokens["obs"])))
            if tokens["tree"]:
                await db.execute(delete(TreeModel).where(TreeModel.id.in_(tokens["tree"])))
            if tokens["combo"]:
                await db.execute(delete(CrossCombinationModel).where(
                    CrossCombinationModel.id.in_(tokens["combo"])))
            if tokens["germ"]:
                await db.execute(delete(BreedingGermplasmModel).where(
                    BreedingGermplasmModel.id.in_(tokens["germ"])))
            await db.execute(delete(TraitModel).where(TraitModel.trait_code == f"tM_{PREFIX}"))
            if tokens["target"]:
                await db.execute(delete(TargetModel).where(TargetModel.id.in_(tokens["target"])))
            await db.commit()
            print(f"[cleanup] 孟德尔 tc 数据已清（check={len(check_ids)}）")
    finally:
        await engine.dispose()


async def _check_isolated(check_id: int) -> None:
    """run1 落库后立即查隔离注册表（必须在 run2 清除之前验）。"""
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            excl = (await db.execute(select(TreePedigreeExclusionModel).where(
                TreePedigreeExclusionModel.is_deleted.is_(False),
                TreePedigreeExclusionModel.tree_id == FIX["T2"]))).scalars().all()
            check("[4] T2 进入隔离注册表（source=mendelian / source_check_id）",
                  len(excl) == 1 and excl[0].source_type == "mendelian"
                  and excl[0].source_check_id == check_id, f"n={len(excl)}")
    finally:
        await engine.dispose()


async def _wipe() -> None:
    """启动前清理本前缀残留（防上次进程被杀/建 fixture 中途失败留下的脏数据）。"""
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            ds_ids = list((await db.execute(select(GenotypingDatasetModel.id).where(
                GenotypingDatasetModel.dataset_name.like(f"{PREFIX}-%")))).scalars())
            tree_ids = list((await db.execute(select(TreeModel.id).where(
                TreeModel.tree_no.like(f"{PREFIX}-%")))).scalars())
            if tree_ids:
                await db.execute(delete(TreePedigreeExclusionModel).where(
                    TreePedigreeExclusionModel.tree_id.in_(tree_ids)))
            check_ids = list((await db.execute(select(MendelianCheckModel.id).where(
                MendelianCheckModel.dataset_id.in_(ds_ids)))).scalars())
            if check_ids:
                await db.execute(delete(MendelianCheckTreeModel).where(
                    MendelianCheckTreeModel.check_id.in_(check_ids)))
                await db.execute(delete(MendelianCheckModel).where(
                    MendelianCheckModel.id.in_(check_ids)))
            if ds_ids:
                sample_ids = list((await db.execute(select(GenotypeSampleModel.id).where(
                    GenotypeSampleModel.dataset_id.in_(ds_ids)))).scalars())
                if sample_ids:
                    await db.execute(delete(GenotypeCallModel).where(
                        GenotypeCallModel.sample_id.in_(sample_ids)))
                    await db.execute(delete(GenotypeSampleModel).where(
                        GenotypeSampleModel.id.in_(sample_ids)))
                await db.execute(delete(GenotypingDatasetModel).where(
                    GenotypingDatasetModel.id.in_(ds_ids)))
            if tree_ids:
                obs_ids = list((await db.execute(select(TraitObservationModel.id).where(
                    TraitObservationModel.tree_id.in_(tree_ids)))).scalars())
                if obs_ids:
                    await db.execute(delete(TraitObservationModel).where(
                        TraitObservationModel.id.in_(obs_ids)))
                await db.execute(delete(TreeModel).where(TreeModel.id.in_(tree_ids)))
            marker_ids = list((await db.execute(select(MarkerModel.id).where(
                MarkerModel.marker_name.like(f"{PREFIX}-%")))).scalars())
            if marker_ids:
                await db.execute(delete(MarkerModel).where(MarkerModel.id.in_(marker_ids)))
            combo_ids = list((await db.execute(select(CrossCombinationModel.id).where(
                CrossCombinationModel.combination_code.like(f"{PREFIX}-%")))).scalars())
            if combo_ids:
                await db.execute(delete(CrossCombinationModel).where(
                    CrossCombinationModel.id.in_(combo_ids)))
            germ_ids = list((await db.execute(select(BreedingGermplasmModel.id).where(
                BreedingGermplasmModel.cultivar_name.like(f"{PREFIX}-%")))).scalars())
            if germ_ids:
                await db.execute(delete(BreedingGermplasmModel).where(
                    BreedingGermplasmModel.id.in_(germ_ids)))
            await db.execute(delete(TraitModel).where(TraitModel.trait_code == f"tM_{PREFIX}"))
            await db.execute(delete(TargetModel).where(
                TargetModel.target_name.like(f"{PREFIX}-%")))
            await db.commit()
            print(f"[preclean] 孟德尔 tc 前缀残留已清（ds={len(ds_ids)} tree={len(tree_ids)}）")
    finally:
        await engine.dispose()


def main_() -> None:
    asyncio.run(_wipe())
    asyncio.run(_build_fixture())
    try:
        with TestClient(create_app()) as client:
            login(client)
            H = auth()
            # ---- [1] POST run ----
            r = client.post("/api/v1/bre/genotype_call/mendelian/check/run",
                            json={"dataset_id": FIX["dataset_id"], "compat_threshold": 0.98},
                            headers=H)
            check("[HTTP] mendelian/check/run 200", r.status_code == 200,
                  f"{r.status_code} {str(r.text)[:150]}")
            b = r.json()
            data = b.get("data") if b.get("code") == 0 else None
            check("[1] 返回 check_id / isolated=1 / cleared=0", data is not None
                  and isinstance(data["check_id"], int) and data["isolated"] == 1
                  and data["cleared"] == 0, f"{data}")
            if data:
                FIX["run1"] = data
                check_ids.append(data["check_id"])
                check("[1] 汇总计数 checked=2 / flagged=1 / no_parents=0 / missing=0",
                      data["trees_checked"] == 2 and data["trees_flagged"] == 1
                      and data["trees_no_parents"] == 0
                      and data["trees_missing_parent_genotype"] == 0,
                      f"{data['trees_checked']}/{data['trees_flagged']}")
                # ---- [4] 隔离注册表：run1 刚落库即验（run2 阈值 0.0 会清除）----
                asyncio.run(_check_isolated(FIX["run1"]["check_id"]))

            # ---- [5] 阈值 0.0 重跑（可逆）----
            r2 = client.post("/api/v1/bre/genotype_call/mendelian/check/run",
                             json={"dataset_id": FIX["dataset_id"], "compat_threshold": 0.0},
                             headers=H)
            b2 = r2.json()
            d2 = b2.get("data") if r2.status_code == 200 and b2.get("code") == 0 else None
            check("[5] 阈值 0.0：flagged=0 / isolated=0 / cleared=1",
                  d2 is not None and d2["trees_flagged"] == 0
                  and d2["isolated"] == 0 and d2["cleared"] == 1, f"{d2}")
            if d2:
                FIX["run2"] = d2
                check_ids.append(d2["check_id"])

            # ---- [6] GET 只读预览零回归 ----
            g = client.get("/api/v1/bre/genotype_call/mendelian/check",
                           params={"dataset_id": FIX["dataset_id"], "compat_threshold": 0.98},
                           headers=H)
            gb = g.json()
            gd = gb.get("data") if g.status_code == 200 and gb.get("code") == 0 else None
            check("[6] GET 只读预览原结构（无 check_id / 纯计算）",
                  gd is not None and "check_id" not in gd
                  and gd["trees_flagged"] == 1 and len(gd["trees"]) == 2
                  and gd["summary"]["incompatible_markers"] == 5,
                  f"flagged={gd and gd['trees_flagged']}")
    finally:
        asyncio.run(_verify_and_cleanup())

    print(f"\n===== 孟德尔 tc 套件：ok={ok} fail={fail} =====")


if __name__ == "__main__":
    main_()
    sys.exit(1 if fail else 0)