Agentswarm/scripts/test-orchestrator.py

#!/usr/bin/env python3
"""Test script for orchestrator functionality."""
import asyncio
import json
import sys
from typing import Optional

try:
    import websockets
    import httpx
except ImportError:
    print("Installing required packages...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets", "httpx"])
    import websockets
    import httpx


ORCHESTRATOR_URL = "http://localhost:8000"
ORCHESTRATOR_WS = "ws://localhost:8000"


async def test_health_check():
    """Test health check endpoint."""
    print("\n=== Testing Health Check ===")
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{ORCHESTRATOR_URL}/health")
        print(f"Status: {response.status_code}")
        print(f"Response: {response.json()}")
        assert response.status_code == 200
        assert response.json()["status"] == "ok"
    print("✅ Health check passed")


async def test_agent_registration():
    """Test agent registration via WebSocket."""
    print("\n=== Testing Agent Registration ===")

    agent_id = "test-agent-1"
    uri = f"{ORCHESTRATOR_WS}/ws/{agent_id}"

    async with websockets.connect(uri) as websocket:
        # Send registration
        await websocket.send(json.dumps({
            "type": "register",
            "capabilities": ["python", "javascript"]
        }))

        # Receive confirmation
        response = await websocket.recv()
        data = json.loads(response)
        print(f"Registration response: {data}")
        assert data["type"] == "registered"
        assert data["agent_id"] == agent_id

        # Verify agent appears in registry
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{ORCHESTRATOR_URL}/agents/{agent_id}")
            agent_data = response.json()
            print(f"Agent data: {agent_data}")
            assert agent_data["agent_id"] == agent_id
            assert agent_data["status"] == "idle"

        print("✅ Agent registration passed")

        # Test heartbeat
        print("\n=== Testing Heartbeat ===")
        await websocket.send(json.dumps({"type": "heartbeat"}))
        response = await websocket.recv()
        data = json.loads(response)
        print(f"Heartbeat response: {data}")
        assert data["type"] == "heartbeat_ack"
        print("✅ Heartbeat passed")


async def test_task_creation():
    """Test task creation and assignment."""
    print("\n=== Testing Task Creation ===")

    async with httpx.AsyncClient() as client:
        # Create task
        response = await client.post(
            f"{ORCHESTRATOR_URL}/tasks",
            json={
                "description": "Test task",
                "context": {"test": "data"},
                "max_retries": 3
            }
        )
        task_data = response.json()
        print(f"Created task: {task_data}")
        assert response.status_code == 200
        assert task_data["status"] == "pending"

        task_id = task_data["task_id"]

        # List tasks
        response = await client.get(f"{ORCHESTRATOR_URL}/tasks")
        tasks = response.json()
        print(f"Total tasks: {len(tasks['tasks'])}")

        print("✅ Task creation passed")
        return task_id


async def test_handoff():
    """Test handoff mechanism between two agents."""
    print("\n=== Testing Handoff Mechanism ===")

    agent1_id = "test-agent-1"
    agent2_id = "test-agent-2"

    # Connect both agents
    uri1 = f"{ORCHESTRATOR_WS}/ws/{agent1_id}"
    uri2 = f"{ORCHESTRATOR_WS}/ws/{agent2_id}"

    async with websockets.connect(uri1) as ws1, websockets.connect(uri2) as ws2:
        # Register agent 1
        await ws1.send(json.dumps({
            "type": "register",
            "capabilities": ["python"]
        }))
        await ws1.recv()

        # Register agent 2
        await ws2.send(json.dumps({
            "type": "register",
            "capabilities": ["javascript"]
        }))
        await ws2.recv()

        # Agent 1 initiates handoff to agent 2
        await ws1.send(json.dumps({
            "type": "handoff",
            "target_agent_id": agent2_id,
            "task_context": {
                "task_id": "test-task-123",
                "description": "Test handoff",
                "files": ["test.py"]
            }
        }))

        # Agent 1 receives handoff confirmation
        response1 = await ws1.recv()
        data1 = json.loads(response1)
        print(f"Agent 1 response: {data1}")
        assert data1["type"] == "handoff_initiated"

        handoff_id = data1["handoff_id"]

        # Agent 2 receives handoff request
        response2 = await ws2.recv()
        data2 = json.loads(response2)
        print(f"Agent 2 response: {data2}")
        assert data2["type"] == "handoff_request"
        assert data2["handoff_id"] == handoff_id

        # Agent 2 accepts handoff
        await ws2.send(json.dumps({
            "type": "handoff_accept",
            "handoff_id": handoff_id
        }))

        response2 = await ws2.recv()
        data2 = json.loads(response2)
        print(f"Agent 2 accept response: {data2}")
        assert data2["type"] == "handoff_accepted"

        # Verify handoff in history
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{ORCHESTRATOR_URL}/handoffs")
            handoffs = response.json()
            print(f"Total handoffs: {len(handoffs['handoffs'])}")
            assert len(handoffs['handoffs']) > 0

        print("✅ Handoff mechanism passed")


async def main():
    """Run all tests."""
    print("Starting orchestrator tests...")
    print("Make sure orchestrator is running on localhost:8000")
    print("Run: kubectl port-forward -n swarm-system svc/orchestrator-service 8000:8000")

    try:
        await test_health_check()
        await test_agent_registration()
        await test_task_creation()
        await test_handoff()

        print("\n" + "="*50)
        print("✅ All tests passed!")
        print("="*50)

    except Exception as e:
        print(f"\n❌ Test failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


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