PlantAI-rag-new/sdk/app/routers/test_kg.py

"""routers/test_kg.py —— kg 路由单元测试。

用假 Neo4jService 替换 services.neo4j.service（monkeypatch），覆盖：
五个端点成功形态、Neo4j 不可用 → 503 UPSTREAM_DOWN、未命中实体 → 404
NOT_FOUND、参数校验（limit ≤ 20、depth 1-2）。
"""
from __future__ import annotations

import httpx
import pytest
from fastapi import FastAPI

from app.routers import kg as kg_router
from app.services import neo4j as neo4j_mod
from app.services.neo4j import Neo4jUnavailable


class FakeService:
    """方法逐项可配的假服务。"""

    def __init__(self, **overrides):
        self.calls: dict[str, list] = {}
        self.overrides = overrides

    async def kg_status(self):
        return self.overrides.get("kg_status", {"ok": True, "active": "x"})

    async def kg_overview(self):
        return self.overrides.get("kg_overview", {"entities_by_type": []})

    async def search_entities(self, q, limit=20):
        self.calls.setdefault("search_entities", []).append((q, limit))
        return self.overrides.get(
            "search_entities",
            [{"id": "cent_ospl14", "name": "OsSPL14", "type": "gene", "support_count": 18}],
        )

    async def entity(self, entity_id):
        self.calls.setdefault("entity", []).append(entity_id)
        return self.overrides.get(
            "entity",
            {"id": "cent_ospl14", "name": "OsSPL14", "type": "gene", "properties": {}},
        )

    async def neighborhood(self, entity_id, depth=1, limit=200):
        self.calls.setdefault("neighborhood", []).append((entity_id, depth, limit))
        return self.overrides.get(
            "neighborhood",
            {
                "center": "cent_ospl14",
                "center_name": "OsSPL14",
                "depth": depth,
                "nodes": [{"id": "cent_ospl14", "label": "OsSPL14", "type": "gene"}],
                "edges": [],
            },
        )


@pytest.fixture
def fake_service(monkeypatch):
    fake = FakeService()
    monkeypatch.setattr(neo4j_mod, "service", fake)
    return fake


@pytest.fixture
def client():
    # 用真实 FastAPI() 包装：自动获得 ExceptionMiddleware + AsyncExitStackMiddleware，
    # 与生产挂载形态一致（生产 main.py 把 kg.router include 进 FastAPI 实例）。
    app = FastAPI()
    app.include_router(kg_router.router)
    return httpx.AsyncClient(
        transport=httpx.ASGITransport(app=app, raise_app_exceptions=False),
        base_url="http://test",
    )

async def test_status_ok(fake_service, client):
    fake_service.overrides["kg_status"] = {
        "ok": True,
        "active": "neo4j-plant-srt-direct-v1-full",
        "schema": "canonical",
        "labels": ["SRT2CanonicalEntity", "SRT2CanonicalAssertion", "SRT2Paper"],
        "relationship_types": ["SRT2_CANONICAL_ASSERTS", "SRT2_SUPPORTED_BY"],
        "entity_count": 3820996,
        "assertion_count": 8246875,
        "name_index": "srt2_canonical_entity_name_norm",
        "fulltext_index": None,
    }
    resp = await client.get("/api/kg/status")
    assert resp.status_code == 200
    assert resp.json()["entity_count"] == 3820996

async def test_status_down_returns_503_envelope(fake_service, client):
    async def boom():
        raise Neo4jUnavailable("all down")

    fake_service.kg_status = boom
    resp = await client.get("/api/kg/status")
    assert resp.status_code == 503
    body = resp.json()
    assert body["ok"] is False
    assert body["error"]["code"] == "UPSTREAM_DOWN"
    assert body["error"]["degraded"] is True

async def test_overview_ok(fake_service, client):
    fake_service.overrides["kg_overview"] = {
        "entities_by_type": [{"type": "gene", "count": 43210}],
        "relations_by_category": [{"type": "characterized_in", "count": 12033}],
    }
    resp = await client.get("/api/kg/overview")
    assert resp.status_code == 200
    body = resp.json()
    assert body["entities_by_type"][0]["type"] == "gene"
    assert body["relations_by_category"][0]["count"] == 12033

async def test_overview_down_returns_503(fake_service, client):
    async def boom():
        raise Neo4jUnavailable("down")

    fake_service.kg_overview = boom
    resp = await client.get("/api/kg/overview")
    assert resp.status_code == 503
    assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"

async def test_entities_ok_and_params_passed(fake_service, client):
    resp = await client.get("/api/kg/entities", params={"q": "OsSPL14", "limit": 5})
    assert resp.status_code == 200
    body = resp.json()
    assert body["query"] == "OsSPL14"
    assert body["results"][0]["id"] == "cent_ospl14"
    assert fake_service.calls["search_entities"] == [("OsSPL14", 5)]

async def test_entities_limit_over_20_rejected(client):
    resp = await client.get("/api/kg/entities", params={"q": "OsSPL14", "limit": 21})
    assert resp.status_code == 422  # FastAPI 参数校验（VALIDATION 语义由 A1 统一处理）

async def test_entities_down_returns_503(fake_service, client):
    async def boom(q, limit=20):
        raise Neo4jUnavailable("down")

    fake_service.search_entities = boom
    resp = await client.get("/api/kg/entities", params={"q": "x"})
    assert resp.status_code == 503
    assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"

async def test_entity_ok(fake_service, client):
    resp = await client.get("/api/kg/entities/cent_ospl14")
    assert resp.status_code == 200
    body = resp.json()
    assert body["id"] == "cent_ospl14"
    assert body["name"] == "OsSPL14"
    assert fake_service.calls["entity"] == ["cent_ospl14"]

async def test_entity_miss_returns_404(fake_service, client):
    fake_service.overrides["entity"] = None
    resp = await client.get("/api/kg/entities/NoSuch")
    assert resp.status_code == 404
    body = resp.json()
    assert body["ok"] is False
    assert body["error"]["code"] == "NOT_FOUND"
    assert body["error"]["degraded"] is False

async def test_entity_down_returns_503(fake_service, client):
    async def boom(entity_id):
        raise Neo4jUnavailable("down")

    fake_service.entity = boom
    resp = await client.get("/api/kg/entities/cent_ospl14")
    assert resp.status_code == 503
    assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"

async def test_neighborhood_ok_and_params(fake_service, client):
    resp = await client.get(
        "/api/kg/entities/cent_ospl14/neighborhood", params={"depth": 2, "limit": 100}
    )
    assert resp.status_code == 200
    body = resp.json()
    assert body["center"] == "cent_ospl14"
    assert body["center_name"] == "OsSPL14"
    assert body["depth"] == 2
    assert fake_service.calls["neighborhood"] == [("cent_ospl14", 2, 100)]

async def test_neighborhood_defaults_depth1_limit50(fake_service, client):
    resp = await client.get("/api/kg/entities/cent_ospl14/neighborhood")
    assert resp.status_code == 200
    assert fake_service.calls["neighborhood"] == [("cent_ospl14", 1, 50)]

async def test_neighborhood_depth_out_of_range_rejected(client):
    resp = await client.get("/api/kg/entities/cent_ospl14/neighborhood", params={"depth": 3})
    assert resp.status_code == 422

async def test_neighborhood_miss_returns_404(fake_service, client):
    fake_service.overrides["neighborhood"] = {
        "center": None, "center_name": None, "depth": 1, "nodes": [], "edges": []
    }
    resp = await client.get("/api/kg/entities/NoSuch/neighborhood")
    assert resp.status_code == 404
    assert resp.json()["error"]["code"] == "NOT_FOUND"

async def test_neighborhood_down_returns_503(fake_service, client):
    async def boom(entity_id, depth=1, limit=200):
        raise Neo4jUnavailable("down")

    fake_service.neighborhood = boom
    resp = await client.get("/api/kg/entities/cent_ospl14/neighborhood")
    assert resp.status_code == 503
    assert resp.json()["error"]["code"] == "UPSTREAM_DOWN"