pipeline works from pi simulation to control output and strategy generation.
This commit is contained in:
Aditya Pulipaka
157
scripts/test_websocket.py
Normal file
157
scripts/test_websocket.py
Normal file
@@ -0,0 +1,157 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Quick test to verify WebSocket control system.
|
||||
Tests the complete flow: Pi → AI → Control Commands
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
|
||||
try:
|
||||
import websockets
|
||||
except ImportError:
|
||||
print("Error: websockets not installed")
|
||||
print("Run: pip install websockets")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
async def test_websocket():
|
||||
"""Test WebSocket connection and control flow."""
|
||||
|
||||
ws_url = "ws://localhost:9000/ws/pi"
|
||||
print(f"Testing WebSocket connection to {ws_url}")
|
||||
print("-" * 60)
|
||||
|
||||
try:
|
||||
async with websockets.connect(ws_url) as websocket:
|
||||
print("✓ WebSocket connected!")
|
||||
|
||||
# 1. Receive welcome message
|
||||
welcome = await websocket.recv()
|
||||
welcome_data = json.loads(welcome)
|
||||
print(f"✓ Welcome message: {welcome_data.get('message')}")
|
||||
|
||||
# 2. Send test telemetry (lap 1)
|
||||
test_payload = {
|
||||
"type": "telemetry",
|
||||
"lap_number": 1,
|
||||
"enriched_telemetry": {
|
||||
"lap": 1,
|
||||
"tire_degradation_rate": 0.15,
|
||||
"pace_trend": "stable",
|
||||
"tire_cliff_risk": 0.05,
|
||||
"optimal_pit_window": [25, 30],
|
||||
"performance_delta": 0.0
|
||||
},
|
||||
"race_context": {
|
||||
"race_info": {
|
||||
"track_name": "Monza",
|
||||
"total_laps": 51,
|
||||
"current_lap": 1,
|
||||
"weather_condition": "Dry",
|
||||
"track_temp_celsius": 28.0
|
||||
},
|
||||
"driver_state": {
|
||||
"driver_name": "Test Driver",
|
||||
"current_position": 5,
|
||||
"current_tire_compound": "medium",
|
||||
"tire_age_laps": 1,
|
||||
"fuel_remaining_percent": 98.0
|
||||
},
|
||||
"competitors": []
|
||||
}
|
||||
}
|
||||
|
||||
print("\n→ Sending lap 1 telemetry...")
|
||||
await websocket.send(json.dumps(test_payload))
|
||||
|
||||
# 3. Wait for response (short timeout for first laps)
|
||||
response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
|
||||
response_data = json.loads(response)
|
||||
|
||||
if response_data.get("type") == "control_command":
|
||||
print("✓ Received control command!")
|
||||
print(f" Brake Bias: {response_data.get('brake_bias')}/10")
|
||||
print(f" Differential Slip: {response_data.get('differential_slip')}/10")
|
||||
print(f" Message: {response_data.get('message', 'N/A')}")
|
||||
else:
|
||||
print(f"✗ Unexpected response: {response_data}")
|
||||
|
||||
# 4. Send two more laps to trigger strategy generation
|
||||
for lap_num in [2, 3]:
|
||||
test_payload["lap_number"] = lap_num
|
||||
test_payload["enriched_telemetry"]["lap"] = lap_num
|
||||
test_payload["race_context"]["race_info"]["current_lap"] = lap_num
|
||||
|
||||
print(f"\n→ Sending lap {lap_num} telemetry...")
|
||||
await websocket.send(json.dumps(test_payload))
|
||||
|
||||
# Lap 3 triggers Gemini, so expect two responses
|
||||
if lap_num == 3:
|
||||
print(f" (lap 3 will trigger strategy generation - may take 10-30s)")
|
||||
|
||||
# First response: immediate acknowledgment
|
||||
response1 = await asyncio.wait_for(websocket.recv(), timeout=5.0)
|
||||
response1_data = json.loads(response1)
|
||||
print(f"✓ Immediate response: {response1_data.get('message', 'Processing...')}")
|
||||
|
||||
# Second response: strategy-based controls
|
||||
print(" Waiting for strategy generation to complete...")
|
||||
response2 = await asyncio.wait_for(websocket.recv(), timeout=45.0)
|
||||
response2_data = json.loads(response2)
|
||||
|
||||
if response2_data.get("type") == "control_command_update":
|
||||
print(f"✓ Lap {lap_num} strategy-based control received!")
|
||||
print(f" Brake Bias: {response2_data.get('brake_bias')}/10")
|
||||
print(f" Differential Slip: {response2_data.get('differential_slip')}/10")
|
||||
|
||||
strategy = response2_data.get('strategy_name')
|
||||
if strategy and strategy != "N/A":
|
||||
print(f" Strategy: {strategy}")
|
||||
print(f" Total Strategies: {response2_data.get('total_strategies')}")
|
||||
print("✓ Strategy generation successful!")
|
||||
else:
|
||||
print(f"✗ Unexpected response: {response2_data}")
|
||||
else:
|
||||
# Laps 1-2: just one response
|
||||
response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
|
||||
response_data = json.loads(response)
|
||||
|
||||
if response_data.get("type") == "control_command":
|
||||
print(f"✓ Lap {lap_num} control command received!")
|
||||
print(f" Brake Bias: {response_data.get('brake_bias')}/10")
|
||||
print(f" Differential Slip: {response_data.get('differential_slip')}/10")
|
||||
|
||||
# 5. Disconnect
|
||||
print("\n→ Sending disconnect...")
|
||||
await websocket.send(json.dumps({"type": "disconnect"}))
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✓ ALL TESTS PASSED!")
|
||||
print("=" * 60)
|
||||
print("\nWebSocket control system is working correctly.")
|
||||
print("Ready to run: python scripts/simulate_pi_websocket.py")
|
||||
|
||||
except websockets.exceptions.WebSocketException as e:
|
||||
print(f"\n✗ WebSocket error: {e}")
|
||||
print("\nMake sure the AI Intelligence Layer is running:")
|
||||
print(" cd ai_intelligence_layer && python main.py")
|
||||
sys.exit(1)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
print("\n✗ Timeout waiting for response")
|
||||
print("AI layer may be processing (Gemini API can be slow)")
|
||||
print("Check ai_intelligence_layer logs for details")
|
||||
sys.exit(1)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n✗ Unexpected error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("WebSocket Control System Test")
|
||||
print("=" * 60)
|
||||
asyncio.run(test_websocket())
|
||||
Reference in New Issue
Block a user