40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
"""Reusable test data and setup helpers."""
|
|
|
|
import httpx
|
|
|
|
from tests.endpoints import Endpoints as E
|
|
|
|
|
|
async def set_scenario(client: httpx.AsyncClient, scenario_id: str) -> dict:
|
|
res = await client.put(E.SCENARIO_ACTIVE, json={"scenario_id": scenario_id})
|
|
return res.json()
|
|
|
|
|
|
async def get_flights(client: httpx.AsyncClient) -> list[dict]:
|
|
return (await client.get(E.DATA_FLIGHTS)).json()
|
|
|
|
|
|
async def get_disrupted_flights(client: httpx.AsyncClient) -> list[dict]:
|
|
return [f for f in await get_flights(client) if f["status"] != "ON_TIME"]
|
|
|
|
|
|
async def get_first_disrupted_flight_id(client: httpx.AsyncClient) -> str:
|
|
flights = await get_disrupted_flights(client)
|
|
assert len(flights) > 0, "No disrupted flights in current scenario"
|
|
return flights[0]["flight_id"]
|
|
|
|
|
|
async def trigger_fce(client: httpx.AsyncClient, flight_id: str) -> str:
|
|
res = await client.post(E.AGENT_FCE, json={"flight_id": flight_id})
|
|
data = res.json()
|
|
assert "run_id" in data, f"FCE trigger failed: {data}"
|
|
return data["run_id"]
|
|
|
|
|
|
async def trigger_handover(client: httpx.AsyncClient, hubs: list[str] | None = None) -> str:
|
|
body = {"hubs": hubs} if hubs else {}
|
|
res = await client.post(E.AGENT_HANDOVER, json=body)
|
|
data = res.json()
|
|
assert "run_id" in data, f"Handover trigger failed: {data}"
|
|
return data["run_id"]
|