taiji-AI-PAD/test_agent_manager_client.py

#!/usr/bin/env python3
"""
Agent Manager 客户端测试脚本

测试 MCP-Server 与 Agent Manager 的集成
"""

import asyncio
import sys
import os

# 添加项目路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'services/mcp-server'))

from app.agent_manager_client import (
    AgentManagerClient,
    AgentConfig,
    AgentManagerError,
    get_agent_manager_client,
)


async def test_health_check(client: AgentManagerClient):
    """测试健康检查"""
    print("\n=== 测试健康检查 ===")
    try:
        result = await client.health_check()
        print(f"✅ 健康检查成功: {result}")
        return True
    except AgentManagerError as e:
        print(f"❌ 健康检查失败: {e.message}")
        return False


async def test_list_templates(client: AgentManagerClient):
    """测试获取所有模板"""
    print("\n=== 测试获取所有模板 ===")
    try:
        templates = await client.list_templates()
        print(f"✅ 获取模板成功，共 {len(templates)} 个模板:")
        for t in templates:
            print(f"   - {t.template} (port: {t.port})")
        return True
    except AgentManagerError as e:
        print(f"❌ 获取模板失败: {e.message}")
        return False


async def test_list_platform_templates(client: AgentManagerClient):
    """测试获取平台模板"""
    print("\n=== 测试获取平台模板 ===")
    try:
        templates = await client.list_platform_templates()
        print(f"✅ 获取平台模板成功，共 {len(templates)} 个模板:")
        for t in templates:
            print(f"   - {t.template} (type: {t.template_type})")
        return True
    except AgentManagerError as e:
        print(f"❌ 获取平台模板失败: {e.message}")
        return False


async def test_list_custom_templates(client: AgentManagerClient):
    """测试获取自定义模板"""
    print("\n=== 测试获取自定义模板 ===")
    try:
        templates = await client.list_custom_templates()
        print(f"✅ 获取自定义模板成功，共 {len(templates)} 个模板:")
        for t in templates:
            env_info = t.env_info
            required = env_info.get("required", {})
            print(f"   - {t.template}")
            if required:
                print(f"     必需环境变量: {list(required.keys())}")
        return True
    except AgentManagerError as e:
        print(f"❌ 获取自定义模板失败: {e.message}")
        return False


async def test_list_agents(client: AgentManagerClient):
    """测试获取所有 Agent"""
    print("\n=== 测试获取所有 Agent ===")
    try:
        result = await client.list_agents()
        print(f"✅ 获取 Agent 列表成功，共 {result.count} 个 Agent:")
        for agent in result.agents[:5]:  # 只显示前5个
            print(f"   - {agent.get('name')} (status: {agent.get('status')})")
        if result.count > 5:
            print(f"   ... 还有 {result.count - 5} 个")
        return True
    except AgentManagerError as e:
        print(f"❌ 获取 Agent 列表失败: {e.message}")
        return False


async def test_create_and_delete_agent(client: AgentManagerClient):
    """测试创建和删除 Agent"""
    print("\n=== 测试创建和删除 Agent ===")
    
    agent_name = "test-echo-agent"
    
    # 创建 Agent
    try:
        config = AgentConfig(
            user_id="test-user-001",
            cpu_request="100m",
            memory_request="128Mi"
        )
        
        print(f"正在创建 Agent: {agent_name}")
        result = await client.create_agent(
            name=agent_name,
            template="echo_agent",
            config=config
        )
        print(f"✅ 创建 Agent 成功:")
        print(f"   - name: {result.name}")
        print(f"   - namespace: {result.namespace}")
        print(f"   - status: {result.status}")
        print(f"   - pod_id: {result.pod_id}")
        
    except AgentManagerError as e:
        if e.status_code == 409:
            print(f"⚠️ Agent 已存在，跳过创建")
        else:
            print(f"❌ 创建 Agent 失败: {e.message}")
            return False
    
    # 获取状态
    try:
        print(f"\n正在获取 Agent 状态: {agent_name}")
        status = await client.get_agent_status(agent_name)
        print(f"✅ 获取状态成功:")
        print(f"   - status: {status.status}")
        print(f"   - pod_ip: {status.pod_ip}")
        print(f"   - node_name: {status.node_name}")
        
    except AgentManagerError as e:
        print(f"❌ 获取状态失败: {e.message}")
    
    # 获取资源使用
    try:
        print(f"\n正在获取 Agent 资源使用: {agent_name}")
        metrics = await client.get_agent_metrics(agent_name)
        print(f"✅ 获取资源使用成功:")
        print(f"   - resources: {metrics.resources}")
        
    except AgentManagerError as e:
        print(f"❌ 获取资源使用失败: {e.message}")
    
    # 删除 Agent
    try:
        print(f"\n正在删除 Agent: {agent_name}")
        result = await client.delete_agent(agent_name)
        print(f"✅ 删除 Agent 成功: {result}")
        
    except AgentManagerError as e:
        print(f"❌ 删除 Agent 失败: {e.message}")
        return False
    
    return True


async def test_not_implemented_methods(client: AgentManagerClient):
    """测试未实现的方法"""
    print("\n=== 测试未实现的方法 ===")
    
    methods = [
        ("scale_agent", lambda: client.scale_agent("test", replicas=2)),
        ("get_agent_logs", lambda: client.get_agent_logs("test")),
        ("restart_agent", lambda: client.restart_agent("test")),
        ("get_resource_stats", lambda: client.get_resource_stats()),
        ("get_user_resources", lambda: client.get_user_resources("test-user")),
        ("get_channel_resources", lambda: client.get_channel_resources("test-channel")),
    ]
    
    for name, method in methods:
        try:
            await method()
            print(f"❌ {name} 应该抛出 NotImplementedError")
        except NotImplementedError as e:
            print(f"✅ {name} 正确抛出 NotImplementedError")
        except Exception as e:
            print(f"⚠️ {name} 抛出其他异常: {e}")
    
    return True


async def main():
    """主测试函数"""
    print("=" * 60)
    print("Agent Manager 客户端测试")
    print("=" * 60)
    
    # 获取 Agent Manager URL
    agent_manager_url = os.getenv("AGENT_MANAGER_URL", "http://localhost:8000")
    print(f"\nAgent Manager URL: {agent_manager_url}")
    
    # 创建客户端
    client = AgentManagerClient(base_url=agent_manager_url)
    
    try:
        # 运行测试
        results = []
        
        # 基础测试
        results.append(("健康检查", await test_health_check(client)))
        
        if results[-1][1]:  # 如果健康检查成功，继续其他测试
            results.append(("获取所有模板", await test_list_templates(client)))
            results.append(("获取平台模板", await test_list_platform_templates(client)))
            results.append(("获取自定义模板", await test_list_custom_templates(client)))
            results.append(("获取 Agent 列表", await test_list_agents(client)))
            results.append(("创建和删除 Agent", await test_create_and_delete_agent(client)))
        
        # 测试未实现的方法
        results.append(("未实现方法检查", await test_not_implemented_methods(client)))
        
        # 打印结果汇总
        print("\n" + "=" * 60)
        print("测试结果汇总")
        print("=" * 60)
        
        passed = 0
        failed = 0
        for name, result in results:
            status = "✅ 通过" if result else "❌ 失败"
            print(f"  {name}: {status}")
            if result:
                passed += 1
            else:
                failed += 1
        
        print(f"\n总计: {passed} 通过, {failed} 失败")
        
    finally:
        await client.close()


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