LabelingSystem/backend/scripts/test_external_api.py

#!/usr/bin/env python3
"""
外部API测试脚本

功能：
1. 测试5种标注类型的项目创建
2. 测试进度查询接口
3. 测试数据导出接口
4. 验证标签传入功能

使用方式：
    python scripts/test_external_api.py --base-url http://localhost:8003 --token <admin_token>
    
    或使用环境变量：
    export API_BASE_URL=http://localhost:8003
    export ADMIN_TOKEN=your_token_here
    python scripts/test_external_api.py
"""
import argparse
import os
import sys
import json
import requests
from typing import Optional, Dict, Any, List
from dataclasses import dataclass
from datetime import datetime


@dataclass
class TestResult:
    """测试结果"""
    name: str
    success: bool
    message: str
    details: Optional[Dict] = None


class ExternalAPITester:
    """外部API测试器"""
    
    def __init__(self, base_url: str, token: str):
        self.base_url = base_url.rstrip('/')
        self.token = token
        self.headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }
        self.results: List[TestResult] = []
        self.created_projects: List[str] = []
    
    def log(self, message: str):
        """打印日志"""
        print(f"[{datetime.now().strftime('%H:%M:%S')}] {message}")
    
    def add_result(self, name: str, success: bool, message: str, details: Optional[Dict] = None):
        """添加测试结果"""
        result = TestResult(name, success, message, details)
        self.results.append(result)
        status = "✓" if success else "✗"
        self.log(f"{status} {name}: {message}")
        if details and not success:
            self.log(f"  详情: {json.dumps(details, ensure_ascii=False, indent=2)}")
    
    def test_create_project(
        self,
        name: str,
        task_type: str,
        data: List[Dict],
        tags: Optional[List[Dict]] = None,
        description: str = ""
    ) -> Optional[str]:
        """
        测试创建项目
        
        Returns:
            str: 项目ID，失败返回None
        """
        test_name = f"创建{task_type}项目"
        
        payload = {
            "name": name,
            "description": description,
            "task_type": task_type,
            "data": data,
            "external_id": f"test_{task_type}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        }
        
        if tags:
            payload["tags"] = tags
        
        try:
            response = requests.post(
                f"{self.base_url}/api/external/projects/init",
                json=payload,
                headers=self.headers,
                timeout=30
            )
            
            if response.status_code == 201:
                result = response.json()
                project_id = result.get("project_id")
                self.created_projects.append(project_id)
                
                # 验证标签是否正确应用
                if tags:
                    config = result.get("config", "")
                    tags_found = all(tag["tag"] in config for tag in tags)
                    if tags_found:
                        self.add_result(
                            test_name,
                            True,
                            f"项目创建成功，ID: {project_id}，标签已正确应用",
                            {"project_id": project_id, "task_count": result.get("task_count")}
                        )
                    else:
                        self.add_result(
                            test_name,
                            False,
                            f"项目创建成功但标签未正确应用",
                            {"project_id": project_id, "config": config[:200]}
                        )
                else:
                    self.add_result(
                        test_name,
                        True,
                        f"项目创建成功，ID: {project_id}",
                        {"project_id": project_id, "task_count": result.get("task_count")}
                    )
                return project_id
            else:
                self.add_result(
                    test_name,
                    False,
                    f"HTTP {response.status_code}",
                    {"response": response.text[:500]}
                )
                return None
                
        except Exception as e:
            self.add_result(test_name, False, f"请求异常: {str(e)}")
            return None
    
    def test_get_progress(self, project_id: str) -> bool:
        """测试获取项目进度"""
        test_name = f"查询项目进度 ({project_id[:20]}...)"
        
        try:
            response = requests.get(
                f"{self.base_url}/api/external/projects/{project_id}/progress",
                headers=self.headers,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                self.add_result(
                    test_name,
                    True,
                    f"进度查询成功，完成率: {result.get('completion_percentage', 0)}%",
                    {
                        "total_tasks": result.get("total_tasks"),
                        "completed_tasks": result.get("completed_tasks"),
                        "status": result.get("status")
                    }
                )
                return True
            else:
                self.add_result(
                    test_name,
                    False,
                    f"HTTP {response.status_code}",
                    {"response": response.text[:500]}
                )
                return False
                
        except Exception as e:
            self.add_result(test_name, False, f"请求异常: {str(e)}")
            return False
    
    def test_export_project(self, project_id: str, format: str = "json") -> bool:
        """测试导出项目数据"""
        test_name = f"导出项目数据 ({project_id[:20]}..., {format}格式)"
        
        try:
            response = requests.post(
                f"{self.base_url}/api/external/projects/{project_id}/export",
                json={"format": format, "completed_only": False},
                headers=self.headers,
                timeout=60
            )
            
            if response.status_code == 200:
                result = response.json()
                self.add_result(
                    test_name,
                    True,
                    f"导出成功，共{result.get('total_exported', 0)}条数据",
                    {
                        "file_url": result.get("file_url"),
                        "file_name": result.get("file_name"),
                        "file_size": result.get("file_size")
                    }
                )
                return True
            else:
                self.add_result(
                    test_name,
                    False,
                    f"HTTP {response.status_code}",
                    {"response": response.text[:500]}
                )
                return False
                
        except Exception as e:
            self.add_result(test_name, False, f"请求异常: {str(e)}")
            return False
    
    def run_all_tests(self):
        """运行所有测试"""
        self.log("=" * 60)
        self.log("开始外部API测试")
        self.log("=" * 60)
        self.log(f"API地址: {self.base_url}")
        self.log("")
        
        # 测试数据
        text_data = [
            {"id": "text_1", "content": "这个产品非常好用，推荐购买！"},
            {"id": "text_2", "content": "质量太差了，不值这个价格"},
            {"id": "text_3", "content": "一般般，没什么特别的"}
        ]
        
        image_data = [
            {"id": "img_1", "content": "https://picsum.photos/id/40/800/600"},
            {"id": "img_2", "content": "https://picsum.photos/id/40/800/600"},
            {"id": "img_3", "content": "https://picsum.photos/id/40/800/600"}
        ]
        
        # 标签定义
        sentiment_tags = [
            {"tag": "正面", "color": "#4CAF50"},
            {"tag": "负面", "color": "#F44336"},
            {"tag": "中性", "color": "#9E9E9E"}
        ]
        
        animal_tags = [
            {"tag": "猫猫", "color": "#FF9800"},
            {"tag": "狗狗", "color": "#2196F3"},
            {"tag": "其他"}  # 不指定颜色，测试自动生成
        ]
        
        object_tags = [
            {"tag": "人物", "color": "#E91E63"},
            {"tag": "车辆", "color": "#00BCD4"},
            {"tag": "建筑", "color": "#795548"}
        ]
        
        ner_tags = [
            {"tag": "人名", "color": "#9C27B0"},
            {"tag": "地名", "color": "#3F51B5"},
            {"tag": "组织", "color": "#009688"}
        ]
        
        polygon_tags = [
            {"tag": "区域A", "color": "#FF5722"},
            {"tag": "区域B", "color": "#607D8B"},
            {"tag": "区域C"}  # 不指定颜色
        ]
        
        # 1. 测试文本分类项目
        self.log("\n--- 测试1: 文本分类项目 ---")
        project_id = self.test_create_project(
            name="测试-文本分类项目",
            task_type="text_classification",
            data=text_data,
            tags=sentiment_tags,
            description="情感分析测试项目"
        )
        if project_id:
            self.test_get_progress(project_id)
            self.test_export_project(project_id, "json")
        
        # 2. 测试图像分类项目
        self.log("\n--- 测试2: 图像分类项目 ---")
        project_id = self.test_create_project(
            name="测试-图像分类项目",
            task_type="image_classification",
            data=image_data,
            tags=animal_tags,
            description="动物分类测试项目"
        )
        if project_id:
            self.test_get_progress(project_id)
            self.test_export_project(project_id, "csv")
        
        # 3. 测试目标检测项目
        self.log("\n--- 测试3: 目标检测项目 ---")
        project_id = self.test_create_project(
            name="测试-目标检测项目",
            task_type="object_detection",
            data=image_data,
            tags=object_tags,
            description="目标检测测试项目"
        )
        if project_id:
            self.test_get_progress(project_id)
            self.test_export_project(project_id, "yolo")
            self.test_export_project(project_id, "coco")
        
        # 4. 测试NER项目
        self.log("\n--- 测试4: 命名实体识别项目 ---")
        ner_data = [
            {"id": "ner_1", "content": "张三在北京的阿里巴巴公司工作"},
            {"id": "ner_2", "content": "李四去了上海参加腾讯的面试"},
            {"id": "ner_3", "content": "王五是华为深圳总部的工程师"}
        ]
        project_id = self.test_create_project(
            name="测试-NER项目",
            task_type="ner",
            data=ner_data,
            tags=ner_tags,
            description="命名实体识别测试项目"
        )
        if project_id:
            self.test_get_progress(project_id)
            self.test_export_project(project_id, "json")
        
        # 5. 测试多边形标注项目
        self.log("\n--- 测试5: 多边形标注项目 ---")
        project_id = self.test_create_project(
            name="测试-多边形标注项目",
            task_type="polygon",
            data=image_data,
            tags=polygon_tags,
            description="多边形区域标注测试项目"
        )
        if project_id:
            self.test_get_progress(project_id)
            self.test_export_project(project_id, "coco")
        
        # 6. 测试不带标签的项目（验证默认行为）
        self.log("\n--- 测试6: 不带标签的项目 ---")
        project_id = self.test_create_project(
            name="测试-无标签项目",
            task_type="text_classification",
            data=text_data,
            tags=None,  # 不传标签
            description="测试默认配置"
        )
        if project_id:
            self.test_get_progress(project_id)
        
        # 输出测试结果汇总
        self.print_summary()
    
    def print_summary(self):
        """打印测试结果汇总"""
        self.log("\n" + "=" * 60)
        self.log("测试结果汇总")
        self.log("=" * 60)
        
        total = len(self.results)
        passed = sum(1 for r in self.results if r.success)
        failed = total - passed
        
        self.log(f"总计: {total} 个测试")
        self.log(f"通过: {passed} 个")
        self.log(f"失败: {failed} 个")
        self.log(f"通过率: {(passed/total*100):.1f}%" if total > 0 else "N/A")
        
        if failed > 0:
            self.log("\n失败的测试:")
            for r in self.results:
                if not r.success:
                    self.log(f"  ✗ {r.name}: {r.message}")
        
        self.log("\n创建的项目ID:")
        for pid in self.created_projects:
            self.log(f"  - {pid}")
        
        self.log("\n" + "=" * 60)
        
        # 返回退出码
        return 0 if failed == 0 else 1


def main():
    parser = argparse.ArgumentParser(description="外部API测试脚本")
    parser.add_argument(
        "--base-url",
        default=os.environ.get("API_BASE_URL", "http://localhost:8003"),
        help="API基础URL (默认: http://localhost:8003)"
    )
    parser.add_argument(
        "--token",
        default=os.environ.get("ADMIN_TOKEN", ""),
        help="管理员Token"
    )
    
    args = parser.parse_args()
    
    if not args.token:
        print("错误: 请提供管理员Token")
        print("使用方式:")
        print("  python scripts/test_external_api.py --token <your_token>")
        print("  或设置环境变量 ADMIN_TOKEN")
        sys.exit(1)
    
    tester = ExternalAPITester(args.base_url, args.token)
    exit_code = tester.run_all_tests()
    sys.exit(exit_code)


if __name__ == "__main__":
    main()