InterfaceAutoTest/core/models.py

#!/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}")