test_new_configs.py

import requests
import json
import asyncio
import websockets

def test_api():
    print("Testing API...")
    url = "http://localhost:8000/api/run"
    
    amm_config = {
        "id": "amm_1",
        "category": "amm",
        "version": "1.0",
        "core_model": {
            "type": "constant_sum",
            "pricing_equation": "constant_sum",
            "invariant": "x + y = k",
            "parameters": {}
        },
        "behavioral_parameters": {
            "swap_fee_percent": 0.5,
            "slippage_curve": "linear",
            "arbitrage_enabled": True,
            "arbitrage_reaction_speed_blocks": 1
        },
        "risk_parameters": {}
    }

    liquidity_config = {
        "id": "lpool_1",
        "category": "liquidity",
        "version": "1.0",
        "core_model": {
            "initial_tvl": 1000000.0,
            "liquidity_distribution": {
                "whales_percent": 80.0,
                "retail_percent": 20.0
            }
        },
        "behavioral_parameters": {
            "withdrawal_trigger_price_drop_percent": 2.0,
            "withdrawal_fraction_on_trigger": 0.8,
            "withdrawal_speed": "instant",
            "volatility_threshold": 5.0,
            "reentry_enabled": False,
            "reentry_delay_blocks": 0
        },
        "risk_parameters": {}
    }

    data = {
        "steps": 20,
        "prompt": "An aggressive whale attacks the peg.",
        "amm_config": json.dumps(amm_config),
        "liquidity_config": json.dumps(liquidity_config)
    }

    response = requests.post(url, params=data)
    if response.status_code != 200:
        print(f"Failed to start run: {response.text}")
        return

    result = response.json()
    run_id = result["run_id"]
    print(f"Started run: {run_id}")
    
    asyncio.run(connect_ws(run_id))

async def connect_ws(run_id):
    uri = f"ws://localhost:8000/ws/{run_id}"
    try:
        async with websockets.connect(uri) as websocket:
            count = 0
            while count < 5: # Just watch the first 5 steps
                msg = await websocket.recv()
                data = json.loads(msg)
                if data.get("type") == "DONE":
                    print("Simulation finished")
                    break
                print(f"Step {data.get('step')}: Pool USDC = {data.get('pool_usdc')}, AMM Price = {data.get('amm_price')}, Oracle = {data.get('oracle_price')}")
                if data.get("type") == "ERROR":
                    print("ERROR:", data.get("msg"))
                    break
                count += 1
            await websocket.close()
    except Exception as e:
        print(f"WebSocket error: {e}")

if __name__ == "__main__":
    test_api()