#!/usr/bin/env python
import asyncio
import websockets
import json

async def test_energy_refill():
    uri = "ws://localhost:8001"
    async with websockets.connect(uri) as websocket:
        # Connect
        await websocket.send(json.dumps({"type": "connect", "userID": "test_user_123"}))
        response = await websocket.recv()
        print(f"Connected: {response}")

        # Get tiers
        await websocket.send(json.dumps({"type": "getEnergyRefillTiers"}))
        tiers_response = await websocket.recv()
        print(f"Tiers: {tiers_response}")

        # Purchase small refill
        await websocket.send(json.dumps({
            "type": "purchaseEnergyRefill",
            "message": {"refillType": "small"}
        }))
        purchase_response = await websocket.recv()
        print(f"Purchase: {purchase_response}")

asyncio.run(test_energy_refill())