test_ws.py

import asyncio
import json
import websockets
import httpx

async def run_scenario_test():
    # 1. Trigger run
    async with httpx.AsyncClient() as client:
        resp = await client.post(
            "http://127.0.0.1:8000/api/run",
            data={"scenario_id": "dex_arbitrage", "steps": 5}
        )
        print("Run created:", resp.json())
        run_id = resp.json() ["run_id"]
        
    # 2. Connect to WS
    uri = f"ws://127.0.0.1:8000/ws/{run_id}"
    print(f"Connecting to {uri}...")
    
    try:    
        async with websockets.connect(uri) as ws:
            async for msg in ws:
                data = json.loads(msg)
                print(f"[{data.get('type')}] Step {data.get('step', 'done')}")
                
                if data.get("type") == "DONE":
                    print("Summary:", json.dumps(data.get("summary"), indent=2))
                    break
    except websockets.exceptions.ConnectionClosed:
        print("WebSocket connection closed cleanly by server.")

if __name__ == "__main__":
    asyncio.run(run_scenario_test())