cma-management/backend/scripts/verify_customer_concentration_rule_e2e.py

#!/usr/bin/env python3
"""任务B 交付验收脚本 — customer_concentration 客户集中度预警（生产 E2E）

做四件事（可复跑）：
  1. 为 entity1 配置常驻客户集中度规则（KPI=C_CUST_CONCENTRATION, source=auto → zxt 客户收入）
  2. 打印真实数据下的单一客户占比（zxt_masterbill 近12个月）作为基线
  3. 触发测试：建带标记的临时规则（source=kpi）+ 带标记的临时 KPIValue（占比 68%）
     → 调 POST /api/cma/alert-rules/run-customer-concentration → 断言 pending 预警生成
  4. 清理临时规则/KPIValue/预警，回查归零

用法：cd /root/cma-management/backend && venv/bin/python3 scripts/verify_customer_concentration_rule_e2e.py
"""
import json
import sys
import urllib.request
import urllib.error

BASE = "http://127.0.0.1:8010"
ADMIN = {"username": "admin", "password": "admin123"}
MARK = "【系统验证临时数据-验证后删除】"
TEMP_TAG = "e2e-verify-temp"
RULE_TYPE = "customer_concentration"
KPI_ID = 23  # C_CUST_CONCENTRATION 客户集中度 (entity1)


def http(method, path, token=None, body=None, timeout=60):
    req = urllib.request.Request(
        BASE + path,
        data=json.dumps(body).encode() if body is not None else None,
        method=method,
    )
    req.add_header("Content-Type", "application/json")
    if token:
        req.add_header("Authorization", f"Bearer {token}")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status, json.loads(r.read().decode() or "{}")
    except urllib.error.HTTPError as e:
        return e.code, {"detail": e.read().decode()[:300]}


def login(entity_id=1):
    st, r = http("POST", "/api/cma/auth/login", body={**ADMIN, "entity_id": entity_id})
    assert st == 200, f"登录失败: {st} {r}"
    return r["token"]


def main():
    sys.path.insert(0, ".")
    from app.database import get_session_local
    from app.models import KPIValue, KPIAlert, KPIDefinition
    from app.api.alert_rules import AlertRule, _customer_revenue_shares

    token = login(1)
    db = get_session_local()()
    created_ids = []
    try:
        # ── 1+2. 常驻规则 + 真实基线 ──────────────────────────────
        print("=== 1. 常驻客户集中度规则（entity1）===")
        st, rules = http("GET", f"/api/cma/alert-rules?rule_type={RULE_TYPE}", token=token)
        assert st == 200, f"列规则失败: {st} {rules}"
        perm = [r for r in rules.get("data", []) if r["kpi_id"] == KPI_ID and r["entity_id"] == 1
                and not (r.get("params") or {}).get("_verify")]
        if perm:
            print(f"  已存在 rule_id={perm[0]['id']}（幂等跳过）")
        else:
            st, r = http("POST", "/api/cma/alert-rules", token=token, body={
                "kpi_id": KPI_ID,
                "rule_type": RULE_TYPE,
                "trigger_on": "actual",
                "enabled": 1,
                "params": {"threshold_pct": 50, "red_ratio": 1.5, "top_n": 1, "source": "auto"},
            })
            assert st == 200, f"建规则失败: {st} {r}"
            print(f"  新建 rule_id={r['data']['id']} params={r['data']['params']}")

        print("=== 2. 真实数据基线（zxt_masterbill 近12个月, entity1）===")
        db.rollback()
        shares, total, rng = _customer_revenue_shares(db, 1, lookback_months=12, top_n=5)
        assert total > 0, "entity1 无客户收入数据 → 基线取数失败"
        print(f"  客户收入合计={total:,.0f}元 ({rng[0]}~{rng[1]}), 客户数={len(shares)}")
        for uid, amt, pct in shares:
            print(f"    客户#{uid}: {amt:,.0f}元 / {pct}%")
        top1 = shares[0][2]
        print(f"  → 单一客户占比 {top1}%（阈值50% → {'触发' if top1 > 50 else '未触发'}）")

        # ── 3. 触发测试（临时规则 + 临时KPI值）────────────────────
        print("=== 3. 触发测试（构造占比超阈值场景）===")
        st, r = http("POST", "/api/cma/alert-rules", token=token, body={
            "kpi_id": KPI_ID,
            "rule_type": RULE_TYPE,
            "trigger_on": "actual",
            "enabled": 1,
            "params": {"threshold_pct": 50, "red_ratio": 1.5, "top_n": 1,
                       "source": "kpi", "_verify": TEMP_TAG},
        })
        assert st == 200, f"建临时规则失败: {st} {r}"
        temp_rule_id = r["data"]["id"]
        created_ids.append(temp_rule_id)
        print(f"  临时规则 rule_id={temp_rule_id} (source=kpi, 阈值50%)")

        db.rollback()
        temp = KPIValue(kpi_id=KPI_ID, entity_id=1, period="2026-08", actual_value=68.0,
                        source_type="manual", data_status="pending", remark=MARK)
        db.add(temp)
        db.commit()
        db.refresh(temp)
        temp_kv_id = temp.id
        print(f"  临时KPIValue id={temp_kv_id} 占比=68.0%")

        st, r = http("POST", "/api/cma/alert-rules/run-customer-concentration", token=token)
        assert st == 200, f"触发失败: {st} {r}"
        print(f"  run-customer-concentration → {r}")

        db.rollback()
        alerts = db.query(KPIAlert).filter(
            KPIAlert.entity_id == 1, KPIAlert.kpi_id == KPI_ID,
            KPIAlert.alert_type == "concentration", KPIAlert.status == "pending",
        ).order_by(KPIAlert.id.desc()).all()
        assert alerts, "未生成客户集中度预警 → 触发测试 FAIL"
        a = alerts[0]
        alert_id = a.id
        print(f"  ✅ 预警生成: id={a.id} level={a.alert_level} type={a.alert_type}")
        print(f"     message={a.alert_message}")
        assert a.alert_level == "yellow", f"68%应为黄灯(50%~75%), 实际={a.alert_level}"

        # ── 4. 清理 ────────────────────────────────────────────
        print("=== 4. 清理临时数据（生产不留假数据）===")
        db.rollback()
        for obj in (db.query(KPIAlert).filter(KPIAlert.id == alert_id).first(),
                    db.query(KPIValue).filter(KPIValue.id == temp_kv_id).first(),
                    db.query(AlertRule).filter(AlertRule.id == temp_rule_id).first()):
            if obj:
                db.delete(obj)
        db.commit()
        left = {
            "alert": db.query(KPIAlert).filter(KPIAlert.id == alert_id).count(),
            "kpi_value": db.query(KPIValue).filter(KPIValue.remark == MARK).count(),
            "temp_rule": db.query(AlertRule).filter(AlertRule.id == temp_rule_id).count(),
        }
        print(f"  回查: {left}")
        assert all(v == 0 for v in left.values())

        # ── 5. 终态 ────────────────────────────────────────────
        print("=== 5. 常驻规则终态 ===")
        st, rules = http("GET", f"/api/cma/alert-rules?rule_type={RULE_TYPE}", token=token)
        for r in rules.get("data", []):
            kpi = db.query(KPIDefinition).filter(KPIDefinition.id == r["kpi_id"]).first()
            print(f"  rule_id={r['id']} entity={r['entity_id']} kpi={kpi.kpi_code if kpi else r['kpi_id']}"
                  f" enabled={r['enabled']} params={r['params']}")
    finally:
        db.close()

    print("\nALL PASS")
    return 0


if __name__ == "__main__":
    sys.exit(main())