taiji-AI-PAD/scripts/test_billing_fix.py

#!/usr/bin/env python3
"""
计费系统修复验证脚本
验证所有P0修复是否正常工作
"""

import requests
import json
import sys
from datetime import datetime
from typing import Dict, Any

# 配置
BASE_URL = "http://localhost:8002"
POSTGRES_CONTAINER = "taiji-postgres"

# 颜色
class Colors:
    GREEN = '\033[0;32m'
    YELLOW = '\033[1;33m'
    RED = '\033[0;31m'
    NC = '\033[0m'

def print_success(msg: str):
    print(f"{Colors.GREEN}✓ {msg}{Colors.NC}")

def print_warning(msg: str):
    print(f"{Colors.YELLOW}⚠ {msg}{Colors.NC}")

def print_error(msg: str):
    print(f"{Colors.RED}✗ {msg}{Colors.NC}")

def test_webhook_health() -> bool:
    """测试1: Webhook健康检查"""
    print("\n测试1: LiteLLM Webhook健康检查")
    try:
        response = requests.get(f"{BASE_URL}/api/v1/billing/litellm-callback/health", timeout=5)
        if response.status_code == 200:
            data = response.json()
            print_success(f"Webhook端点正常: {json.dumps(data)}")
            return True
        else:
            print_error(f"Webhook返回错误状态码: {response.status_code}")
            return False
    except Exception as e:
        print_error(f"Webhook连接失败: {e}")
        return False

def test_database_tables() -> bool:
    """测试2: 数据库表结构验证"""
    print("\n测试2: 验证数据库表结构")
    import subprocess
    
    try:
        # 检查 agent_billing_records 表
        result = subprocess.run(
            ["docker", "exec", POSTGRES_CONTAINER, "psql", "-U", "postgres", "taiji", 
             "-c", "SELECT COUNT(*) FROM agent_billing_records;"],
            capture_output=True, text=True, timeout=10
        )
        if result.returncode == 0:
            print_success("agent_billing_records表存在且可访问")
        else:
            print_error("agent_billing_records表不可访问")
            return False
        
        # 检查 model_billing_records 表
        result = subprocess.run(
            ["docker", "exec", POSTGRES_CONTAINER, "psql", "-U", "postgres", "taiji", 
             "-c", "SELECT COUNT(*) FROM model_billing_records;"],
            capture_output=True, text=True, timeout=10
        )
        if result.returncode == 0:
            print_success("model_billing_records表存在且可访问")
            return True
        else:
            print_error("model_billing_records表不可访问")
            return False
            
    except Exception as e:
        print_error(f"数据库表验证失败: {e}")
        return False

def test_webhook_callback_simulation() -> bool:
    """测试3: 模拟LiteLLM回调"""
    print("\n测试3: 模拟LiteLLM Token计费回调")
    
    callback_data = {
        "id": f"test_call_{datetime.now().timestamp()}",
        "model": "gpt-4",
        "usage": {
            "prompt_tokens": 100,
            "completion_tokens": 50,
            "total_tokens": 150
        },
        "response_cost": 0.0045,
        "status": "success",
        "startTime": datetime.utcnow().isoformat() + "Z",
        "endTime": datetime.utcnow().isoformat() + "Z",
        "response_time": 1.5,
        "metadata": {
            "tenant_id": "test-tenant"
        }
    }
    
    try:
        response = requests.post(
            f"{BASE_URL}/api/v1/billing/litellm-callback",
            json=callback_data,
            timeout=10
        )
        
        if response.status_code == 200:
            data = response.json()
            print_success(f"Webhook回调成功: record_id={data.get('record_id')}")
            return True
        else:
            print_error(f"Webhook回调失败: {response.status_code} - {response.text}")
            return False
            
    except Exception as e:
        print_error(f"Webhook回调测试失败: {e}")
        return False

def test_litellm_config() -> bool:
    """测试4: 验证LiteLLM配置"""
    print("\n测试4: 验证LiteLLM配置")
    
    import yaml
    config_path = "services/model-gateway/config/litellm.yaml"
    
    try:
        with open(config_path, 'r') as f:
            config = yaml.safe_load(f)
        
        # 检查webhook配置
        general = config.get('general_settings', {})
        success_callbacks = general.get('success_callback', [])
        
        if 'webhook' in success_callbacks:
            print_success("LiteLLM success_callback配置正确")
        else:
            print_error("LiteLLM缺少webhook回调配置")
            return False
        
        webhook_url = general.get('webhook_url')
        if webhook_url and 'billing/litellm-callback' in webhook_url:
            print_success(f"Webhook URL配置正确: {webhook_url}")
            return True
        else:
            print_error(f"Webhook URL配置错误: {webhook_url}")
            return False
            
    except Exception as e:
        print_error(f"LiteLLM配置验证失败: {e}")
        return False

def main():
    print("=" * 50)
    print("计费系统修复验证")
    print("=" * 50)
    
    results = []
    
    # 执行测试
    results.append(("Webhook健康检查", test_webhook_health()))
    results.append(("数据库表结构", test_database_tables()))
    results.append(("Webhook回调模拟", test_webhook_callback_simulation()))
    results.append(("LiteLLM配置", test_litellm_config()))
    
    # 汇总结果
    print("\n" + "=" * 50)
    print("测试结果汇总")
    print("=" * 50)
    
    passed = 0
    failed = 0
    
    for test_name, result in results:
        if result:
            print_success(f"{test_name}: 通过")
            passed += 1
        else:
            print_error(f"{test_name}: 失败")
            failed += 1
    
    print(f"\n总计: {passed} 通过, {failed} 失败")
    
    if failed == 0:
        print_success("\n所有测试通过！计费系统修复成功。")
        return 0
    else:
        print_error(f"\n{failed} 个测试失败，请检查问题。")
        return 1

if __name__ == "__main__":
    sys.exit(main())