test_client.py

#!/usr/bin/env python3
"""
测试客户端 - 用于测试Rovo Dev OpenAI API
"""
import asyncio
import json
import sys
from typing import Dict, Any

import httpx


class RovoAPIClient:
    """Rovo API测试客户端"""
    
    def __init__(self, base_url: str = "http://localhost:8000", api_key: str = None):
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key
        self.headers = {
            "Content-Type": "application/json"
        }
        if api_key:
            self.headers["Authorization"] = f"Bearer {api_key}"
    
    async def test_health(self):
        """测试健康检查"""
        print("🔍 测试健康检查...")
        async with httpx.AsyncClient() as client:
            try:
                response = await client.get(f"{self.base_url}/health")
                print(f"✅ 健康检查: {response.status_code}")
                print(f"   响应: {response.json()}")
                return response.status_code == 200
            except Exception as e:
                print(f"❌ 健康检查失败: {e}")
                return False
    
    async def test_models(self):
        """测试模型列表"""
        print("\n🔍 测试模型列表...")
        async with httpx.AsyncClient() as client:
            try:
                response = await client.get(
                    f"{self.base_url}/v1/models",
                    headers=self.headers
                )
                print(f"✅ 模型列表: {response.status_code}")
                if response.status_code == 200:
                    data = response.json()
                    print(f"   可用模型: {[model['id'] for model in data['data']]}")
                else:
                    print(f"   错误: {response.text}")
                return response.status_code == 200
            except Exception as e:
                print(f"❌ 模型列表测试失败: {e}")
                return False
    
    async def test_chat_completion(self, message: str = "你好，请介绍一下你自己"):
        """测试聊天完成"""
        print(f"\n🔍 测试聊天完成: {message}")
        
        payload = {
            "model": "rovo-dev",
            "messages": [
                {"role": "user", "content": message}
            ],
            "stream": False,
            "temperature": 0.3
        }
        
        async with httpx.AsyncClient(timeout=60.0) as client:
            try:
                response = await client.post(
                    f"{self.base_url}/v1/chat/completions",
                    headers=self.headers,
                    json=payload
                )
                print(f"✅ 聊天完成: {response.status_code}")
                
                if response.status_code == 200:
                    data = response.json()
                    content = data['choices'][0]['message']['content']
                    print(f"   响应内容: {content[:200]}...")
                    if 'usage' in data:
                        usage = data['usage']
                        print(f"   Token使用: {usage}")
                else:
                    print(f"   错误: {response.text}")
                
                return response.status_code == 200
            except Exception as e:
                print(f"❌ 聊天完成测试失败: {e}")
                return False
    
    async def test_stream_completion(self, message: str = "请解释什么是人工智能"):
        """测试流式聊天完成"""
        print(f"\n🔍 测试流式聊天完成: {message}")
        
        payload = {
            "model": "rovo-dev",
            "messages": [
                {"role": "user", "content": message}
            ],
            "stream": True,
            "temperature": 0.3
        }
        
        async with httpx.AsyncClient(timeout=60.0) as client:
            try:
                async with client.stream(
                    "POST",
                    f"{self.base_url}/v1/chat/completions",
                    headers=self.headers,
                    json=payload
                ) as response:
                    print(f"✅ 流式响应: {response.status_code}")
                    
                    if response.status_code == 200:
                        print("   流式内容:")
                        async for line in response.aiter_lines():
                            if line.startswith("data: "):
                                data_str = line[6:]  # 移除 "data: " 前缀
                                if data_str.strip() == "[DONE]":
                                    print("   [完成]")
                                    break
                                try:
                                    data = json.loads(data_str)
                                    if 'choices' in data and data['choices']:
                                        delta = data['choices'][0].get('delta', {})
                                        if 'content' in delta:
                                            print(delta['content'], end='', flush=True)
                                except json.JSONDecodeError:
                                    continue
                        print()  # 换行
                    else:
                        print(f"   错误: {await response.aread()}")
                    
                    return response.status_code == 200
            except Exception as e:
                print(f"❌ 流式聊天完成测试失败: {e}")
                return False
    
    async def run_all_tests(self):
        """运行所有测试"""
        print("🚀 开始运行Rovo Dev OpenAI API测试...")
        
        tests = [
            self.test_health(),
            self.test_models(),
            self.test_chat_completion(),
            self.test_stream_completion()
        ]
        
        results = []
        for test in tests:
            result = await test
            results.append(result)
        
        print(f"\n📊 测试结果: {sum(results)}/{len(results)} 通过")
        
        if all(results):
            print("🎉 所有测试通过！")
            return True
        else:
            print("❌ 部分测试失败")
            return False


async def main():
    """主函数"""
    # 从命令行参数获取配置
    base_url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000"
    api_key = sys.argv[2] if len(sys.argv) > 2 else None
    
    if not api_key:
        print("警告: 未提供API密钥，如果服务器需要认证可能会失败")
        print("使用方法: python test_client.py <base_url> <api_key>")
    
    client = RovoAPIClient(base_url, api_key)
    success = await client.run_all_tests()
    
    sys.exit(0 if success else 1)


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