agent_management/test_swarm_api.py

#!/usr/bin/env python3
"""
Test script for Swarm API endpoints.
"""

import requests
import json
import time

BASE_URL = "http://localhost:8000"


def test_create_swarm():
    """Test creating a swarm"""
    print("\n=== Testing Swarm Creation ===")

    payload = {
        "task_description": "实现用户认证模块，包括登录、注册、密码重置功能",
        "project_context": {
            "repo_url": "https://github.com/test/project",
            "branch": "feature/auth",
            "language": "python",
            "framework": "fastapi"
        },
        "agents": [
            {
                "role": "architect",
                "template": "a2a_litellm_agent",
                "model": "gpt-4",
                "capabilities": ["design", "planning"],
                "system_prompt": "你是架构师，负责设计系统架构和API接口",
                "replicas": 1
            },
            {
                "role": "coder",
                "template": "code_manager_agent",
                "model": "gpt-4",
                "capabilities": ["coding", "git"],
                "system_prompt": "你是开发工程师，负责实现代码",
                "replicas": 2
            },
            {
                "role": "reviewer",
                "template": "a2a_litellm_agent",
                "model": "gpt-4",
                "capabilities": ["review", "testing"],
                "system_prompt": "你是代码审查员，负责审查代码质量",
                "replicas": 1
            }
        ],
        "orchestration": {
            "strategy": "sequential",
            "max_iterations": 3,
            "timeout_minutes": 30
        },
        "owner_id": "test_user"
    }

    try:
        response = requests.post(f"{BASE_URL}/api/swarm/create", json=payload)
        response.raise_for_status()
        result = response.json()

        print(f"✅ Swarm created successfully!")
        print(f"   Swarm ID: {result['swarm_id']}")
        print(f"   Status: {result['status']}")
        print(f"   Agents: {len(result['agents'])}")

        for agent in result['agents']:
            print(f"     - {agent['role']}: {agent['agent_id']} ({agent['status']})")

        return result['swarm_id']

    except requests.exceptions.RequestException as e:
        print(f"❌ Failed to create swarm: {e}")
        if hasattr(e.response, 'text'):
            print(f"   Response: {e.response.text}")
        return None


def test_get_swarm_status(swarm_id):
    """Test getting swarm status"""
    print(f"\n=== Testing Swarm Status (ID: {swarm_id}) ===")

    try:
        response = requests.get(f"{BASE_URL}/api/swarm/{swarm_id}/status")
        response.raise_for_status()
        result = response.json()

        print(f"✅ Swarm status retrieved!")
        print(f"   Status: {result['status']}")
        print(f"   Phase: {result.get('phase', 'N/A')}")
        print(f"   Progress: {result['progress']}%")
        print(f"   Total Messages: {result['metrics']['total_messages']}")
        print(f"   Elapsed: {result['metrics']['elapsed_seconds']}s")

        return result

    except requests.exceptions.RequestException as e:
        print(f"❌ Failed to get swarm status: {e}")
        return None


def test_stream_swarm_results(swarm_id, duration=10):
    """Test streaming swarm results"""
    print(f"\n=== Testing Swarm Results Stream (ID: {swarm_id}) ===")
    print(f"Streaming for {duration} seconds...")

    try:
        response = requests.get(
            f"{BASE_URL}/api/swarm/{swarm_id}/results",
            stream=True,
            timeout=duration + 5
        )
        response.raise_for_status()

        start_time = time.time()
        event_count = 0

        for line in response.iter_lines():
            if time.time() - start_time > duration:
                break

            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    event_count += 1
                    data = json.loads(line[6:])
                    event_type = data.get('type', 'unknown')
                    print(f"   Event #{event_count}: {event_type}")

                    if event_type == 'phase_change':
                        print(f"      Phase: {data.get('phase')}")
                    elif event_type == 'agent_message':
                        print(f"      From: {data.get('from_agent_id', 'orchestrator')}")
                        print(f"      To: {data.get('to_agent_id', 'orchestrator')}")
                    elif event_type == 'result':
                        print(f"      Final Status: {data.get('status')}")
                        break

        print(f"✅ Received {event_count} events")

    except requests.exceptions.RequestException as e:
        print(f"❌ Failed to stream results: {e}")


def test_stop_swarm(swarm_id):
    """Test stopping a swarm"""
    print(f"\n=== Testing Swarm Stop (ID: {swarm_id}) ===")

    payload = {
        "reason": "Test completed",
        "cleanup": False  # Don't cleanup for testing
    }

    try:
        response = requests.post(f"{BASE_URL}/api/swarm/{swarm_id}/stop", json=payload)
        response.raise_for_status()
        result = response.json()

        print(f"✅ Swarm stopped successfully!")
        print(f"   Status: {result['status']}")
        print(f"   Stopped at: {result['stopped_at']}")

        return result

    except requests.exceptions.RequestException as e:
        print(f"❌ Failed to stop swarm: {e}")
        return None


def test_get_swarm_logs(swarm_id):
    """Test getting swarm logs"""
    print(f"\n=== Testing Swarm Logs (ID: {swarm_id}) ===")

    try:
        response = requests.get(f"{BASE_URL}/api/swarm/{swarm_id}/logs")
        response.raise_for_status()
        result = response.json()

        print(f"✅ Swarm logs retrieved!")
        print(f"   Agents: {len(result.get('agents', []))}")

        for agent in result.get('agents', []):
            print(f"     - {agent['role']} ({agent['agent_id']})")

        return result

    except requests.exceptions.RequestException as e:
        print(f"❌ Failed to get swarm logs: {e}")
        return None


def main():
    """Run all tests"""
    print("=" * 60)
    print("Swarm API Test Suite")
    print("=" * 60)

    # Test 1: Create swarm
    swarm_id = test_create_swarm()
    if not swarm_id:
        print("\n❌ Cannot continue tests without swarm_id")
        return

    # Wait a bit for initialization
    print("\nWaiting 2 seconds for initialization...")
    time.sleep(2)

    # Test 2: Get status
    test_get_swarm_status(swarm_id)

    # Test 3: Stream results (for 10 seconds)
    test_stream_swarm_results(swarm_id, duration=10)

    # Test 4: Get logs
    test_get_swarm_logs(swarm_id)

    # Test 5: Stop swarm
    test_stop_swarm(swarm_id)

    # Final status check
    print("\nFinal status check...")
    test_get_swarm_status(swarm_id)

    print("\n" + "=" * 60)
    print("Test Suite Completed!")
    print("=" * 60)


if __name__ == "__main__":
    main()