API压测框架-locust

from locust import HttpUser, task, between
import os

class OCRUser(HttpUser):
    wait_time = between(1, 3)
    
    # 定义测试文件集及对应权重 [文件路径, 权重]
    test_files = [
        ("./sample/1.jpg", 1),   # 1启用的测试用例
        ("./sample/1.heif", 1),
        ("./sample/2.jpg", 1),
        ("./sample/2.heif", 1),
        ("./sample/3.jpg", 1),
        ("./sample/3.heif", 1),
        ("./sample/4.jpg", 1),
        ("./sample/4.heif", 1),
        ("./sample/5.jpg", 1),
        ("./sample/5.heif", 1),
    ]

    def on_start(self):
        """检查所有测试文件是否存在"""
        for path, _ in self.test_files:
            if not os.path.exists(path):
                raise FileNotFoundError(f"测试文件 {path} 不存在")

    @task
    def perform_ocr_test(self):
        """参数化测试入口"""
        for file_path, weight in self.test_files:
            if weight <= 0:
                continue
            self._execute_ocr_test(file_path)

    def _execute_ocr_test(self, file_path):
        """通用测试执行逻辑"""
        file_size = os.path.getsize(file_path) // 1024
        file_type = os.path.splitext(file_path)[1][1:].upper()
        
        with self._post_file(file_path, file_type) as res:
            self._validate_response(res, file_path)

    def _post_file(self, file_path, file_type):
        """统一封装文件上传操作"""
        task_name = f"{os.path.basename(file_path)}-{os.path.getsize(file_path)//1024}k-{file_type}"
        return self.client.post(
            "/recognize/document",
            files={"image": (
                os.path.basename(file_path), 
                open(file_path, "rb"),
                self._get_mime_type(file_type)
            )},
            name=task_name,
            catch_response=True
        )

    def _validate_response(self, response, file_path):
        """统一响应验证逻辑"""
        if response.status_code != 200:
            response.failure(f"[{file_path}] 状态码异常: {response.status_code}")
        elif not response.json():
            response.failure(f"[{file_path}] 无效的JSON响应")

    def _get_mime_type(self, file_type):
        """MIME类型映射保持不变"""
        file_type = file_type.lower()
        return {
            "jpg": "image/jpeg",
            "heif": "image/heif"
        }.get(file_type, "application/octet-stream")

("./sample/1.jpg", 1),   # 启用的测试用例
("./sample/1.heif", 1),