#!/usr/bin/env python # coding=utf-8 """ @author: chen wei @Software: PyCharm @contact: t6i888@163.com @file: models.py @date: 2024 2024/9/15 21:14 @desc: 声明yaml用例格式 """ import logging from typing import List, Any, Optional, Union, Annotated import yaml from pydantic import BaseModel, Field, ConfigDict, model_validator, field_validator, AfterValidator from core import settings logger = logging.getLogger(__name__) def smart_cast_int(v: Any) -> Any: if isinstance(v, str) and v.startswith("${") and v.endswith("}"): return v try: return int(v) except (ValueError, TypeError): return v def smart_cast_dict(v: Any) -> Any: """确保字典格式,若是占位符(字符串形式)则放行""" if isinstance(v, str) and v.startswith("${") and v.endswith("}"): return v if isinstance(v, dict) or v is None: return v return v # 也可以根据需求抛出异常 # 使用 Annotated 定义带校验的类型 SmartInt = Annotated[Union[int, str], AfterValidator(smart_cast_int)] SmartDict = Annotated[Union[dict[str, Any], str], AfterValidator(smart_cast_dict)] # --- 基础请求模型 (用于第一种示例:直接请求) --- class RequestModel(BaseModel): method: str = Field(..., description="HTTP 请求方法: get, post, etc.") url: str = Field(..., description="接口路径或完整 URL") params: Optional[SmartDict] = None data: Optional[SmartDict] = None json_body: Optional[Any] = Field(None, alias="json") headers: Optional[SmartDict] = None cookies: Optional[dict[str, str]] = None timeout: SmartInt = Field(default=10) files: Optional[SmartDict] = None model_config = ConfigDict(extra="allow", populate_by_name=True) # 允许扩展 requests 的其他参数 # --- PO 动作模型 (用于第二种示例:业务层反射调用) --- class ApiActionModel(BaseModel): api_class: str = Field(..., alias="class", description="要调用的 API 类名") method: str = Field(..., description="类中的方法名") params: Optional[SmartDict] = Field(default_factory=dict, description="传给方法的参数") model_config = ConfigDict(populate_by_name=True) # --- 新增:断言条目模型 --- class ValidateItem(BaseModel): check: Any = Field(..., description="要检查的字段或表达式") expect: Any = Field(..., description="期望值") assert_method: str = Field(default="equals", alias="assert", description="断言方法") msg: Optional[str] = Field(default="Assertion", description="断言描述") model_config = ConfigDict(populate_by_name=True) # --- 核心用例数据模型 --- class CaseInfo(BaseModel): # 公共元数据 title: str = Field(..., description="用例标题") epic: Optional[str] = None feature: Optional[str] = None story: Optional[str] = None # 核心逻辑分叉:可以是 request 对象,也可以是 api_action 对象 # 根据 WorkflowExecutor 的逻辑,这里设为 Optional,但在具体校验时可以互斥 request: Optional[RequestModel] = None api_action: Optional[ApiActionModel] = None # 后置处理 extract: Optional[dict[str, List[Any]]] = Field( default=None, description="变量提取表达式,格式: {变量名: [来源, 表达式, 索引]}" ) validate_data: Optional[List[Union[ValidateItem, dict[str, Any]]]] = Field( default_factory=list, alias="validate", description="断言信息" ) # 参数化(在 DataLoader 阶段会被拆解,但在初始加载时需要定义) parametrize: Optional[List[List[Any]]] = None model_config = ConfigDict( populate_by_name=True, # 无论是在代码中用 api_class 还是在 YAML 中用 class 赋值,Pydantic 都能正确识别。 arbitrary_types_allowed=True # 允许在模型中使用非 Pydantic 标准类型(如自定义类实例) ) # 核心优化:增加互斥校验 @model_validator(mode='after') def check_action_type(self) -> 'CaseInfo': if not self.request and not self.api_action: raise ValueError("用例必须包含 'request' 或 'api_action' 其中之一") if self.request and self.api_action: raise ValueError("'request' 和 'api_action' 不能同时存在") return self def is_po_mode(self) -> bool: """判断是否为 PO 模式""" return self.api_action is not None if __name__ == '__main__': # 模拟数据 1:标准请求模式 raw_case_1 = { "title": "查询状态信息", "request": { "method": "get", "url": "/api/v1/info", "headers": {"User-Agent": "pytest-ai"} }, "validate": [ {"check": "status_code", "assert": "equals", "expect": 200, "msg": "响应码200"}, {"check": "$.msg", "expect": "Success"} ] } # 模拟数据 2:PO 模式 (反射调用) raw_case_2 = { "title": "用户登录测试", "api_action": { "class": "UserAPI", "method": "login", "params": {"user": "admin", "pwd": "123"} }, "extract": { "token": ["json", "$.data.token", 0] } } print("--- 开始模型校验测试 ---\n") try: # 验证模式 1 case1 = CaseInfo(**raw_case_1) print(f"✅ 模式1 (Request) 校验通过: {case1.title}") print(f" 请求URL: {case1.request.url}") print(f" 第一个断言方法: {case1.validate_data[0].assert_method}\n") # 验证模式 2 case2 = CaseInfo(**raw_case_2) print(f"✅ 模式2 (PO Mode) 校验通过: {case2.title}") print(f" 调用类: {case2.api_action.api_class}") print(f" 提取规则数: {len(case2.extract)}\n") # 验证非法数据(如:既没有 request 也没有 api_action 的情况可以在业务层进一步校验) # 这里演示 Pydantic 自动类型转换 invalid_data = {"title": "错误用例", "request": {"url": "/api"}} # 缺少 method print("--- 预期失败测试 ---") CaseInfo(**invalid_data) except Exception as e: print(f"❌ 预期内的校验失败: \n{e}")