Spaces:
Sleeping
Sleeping
File size: 2,611 Bytes
f927995 f59a149 f927995 f59a149 f927995 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import pytest
import os
from httpx import AsyncClient, ASGITransport
from app.main import app
# Configuration for test environment
BASE_URL = "http://localhost:8080"
GEMINI_KEY_PRESENT = bool(os.getenv("GEMINI_API_KEY"))
@pytest.fixture
async def client():
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url=BASE_URL) as ac:
yield ac
@pytest.mark.asyncio
async def test_healthz_ok(client):
response = await client.get("/healthz")
assert response.status_code == 200
data = response.json()
assert data["status"] == "ok"
assert isinstance(data["gnn_loaded"], bool)
@pytest.mark.asyncio
async def test_graph_returns_nodes(client):
response = await client.get("/api/v1/graph")
assert response.status_code == 200
data = response.json()
assert "nodes" in data
assert len(data["nodes"]) > 20
@pytest.mark.asyncio
async def test_drugs_filter_tier_1(client):
response = await client.get("/api/v1/drugs", params={"tier": 1})
assert response.status_code == 200
data = response.json()
for drug in data["drugs"]:
assert drug["nlem_tier"] == "TIER_1"
@pytest.mark.asyncio
async def test_drug_detail_paracetamol(client):
response = await client.get("/api/v1/drug/paracetamol")
assert response.status_code == 200
data = response.json()
assert data["drug"]["id"] == "paracetamol"
assert "criticality_breakdown" in data
@pytest.mark.asyncio
async def test_alerts_critical(client):
response = await client.get("/api/v1/alerts", params={"severity": "CRITICAL"})
assert response.status_code == 200
data = response.json()
# Check if our known paracetamol alert is present
alert_ids = [a["id"] for a in data["alerts"]]
assert "alert_paracetamol_2024_q1" in alert_ids
@pytest.mark.asyncio
@pytest.mark.skipif(not GEMINI_KEY_PRESENT, reason="Gemini API key missing")
async def test_query_grounded(client):
response = await client.post("/api/v1/query", json={
"question": "Which drugs depend most on Hebei?"
})
assert response.status_code == 200
data = response.json()
assert len(data["answer"]) > 30
assert len(data["citations"]) > 0
@pytest.mark.asyncio
async def test_simulate_hebei_paracetamol(client):
response = await client.post("/api/v1/simulate", json={
"province": "Hebei",
"duration_days": 14,
"severity": "full_shutdown"
})
assert response.status_code == 200
data = response.json()
affected_ids = [d["id"] for d in data["affected_drugs"]]
assert "paracetamol" in affected_ids
|