CNWei
6393414ab2
- 为 core 目录下主要模块 (models, context, creator, base_api, exchange, executor) 添加了详细的类和方法 Docstring。 - 新增 docs/架构改进.md 文件。
210 lines
7.4 KiB
Python
210 lines
7.4 KiB
Python
#!/usr/bin/env python
|
||
# coding=utf-8
|
||
|
||
"""
|
||
@desc: 变量交换器，用于数据替换和提取
|
||
"""
|
||
|
||
import logging
|
||
import re
|
||
from typing import Any, Union, TypeVar
|
||
|
||
import jsonpath
|
||
from lxml import etree
|
||
|
||
from core.models import RawSchema
|
||
from core.settings import EXTRACT_CACHE
|
||
from core.templates import Template
|
||
from commons.file_processors.yaml_processor import YamlProcessor
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 定义泛型，用于保持返回类型一致
|
||
T = TypeVar("T", bound=Union[dict, list, str, Any])
|
||
|
||
|
||
class Exchange:
|
||
"""
|
||
变量交换器类。
|
||
负责管理全局变量缓存，核心职能包括：
|
||
1. Extract: 从响应结果中提取变量。
|
||
2. Replace: 将数据中的变量占位符替换为实际值。
|
||
"""
|
||
def __init__(self, variable_cache: dict[str, Any]):
|
||
"""
|
||
初始化交换器。
|
||
|
||
Args:
|
||
variable_cache: 初始变量缓存字典（引用传递，修改会影响源数据）。
|
||
"""
|
||
self._cache = variable_cache
|
||
# 匹配标准变量 ${var}，排除函数调用 ${func()}
|
||
self.var_only_pattern = re.compile(r"^\$\{([a-zA-Z_]\w*)}$")
|
||
|
||
@property
|
||
def global_vars(self) -> dict:
|
||
"""获取当前全局变量缓存。"""
|
||
return self._cache
|
||
|
||
@global_vars.setter
|
||
def global_vars(self, global_vars: dict) -> None:
|
||
"""设置全局变量缓存（通常用于上下文切换，如 ChainMap 合并）。"""
|
||
self._cache = global_vars
|
||
|
||
def extract(self, resp: Any, var_name: str, attr: str, expr: str, index: int = 0):
|
||
"""
|
||
从响应中提取数据并更新到缓存及文件
|
||
:param resp: Response 对象
|
||
:param var_name: 变量名
|
||
:param attr: 属性名 (json, text, headers 等)
|
||
:param expr: 提取表达式 ($.jsonpath, //xpath, regex)
|
||
:param index: 索引
|
||
"""
|
||
|
||
try:
|
||
# 兼容处理 resp.json
|
||
target_data = getattr(resp, attr, None)
|
||
if attr == "json":
|
||
try:
|
||
target_data = resp.json()
|
||
except Exception:
|
||
target_data = {"msg": "not json data"}
|
||
|
||
if target_data is None:
|
||
logger.warning(f"提取失败: 响应对象中不存在属性 '{attr}'")
|
||
return
|
||
|
||
value = None
|
||
|
||
if expr.startswith("$"): # JSONPath
|
||
res = jsonpath.jsonpath(target_data, expr)
|
||
if res: value = res[index]
|
||
elif expr.startswith("/") or expr.startswith("./"): # XPath 模式
|
||
html_content = getattr(resp, "text", "") # 使用 getattr 防护
|
||
if not html_content:
|
||
logger.warning("XPath 提取失败：响应文本为空")
|
||
return
|
||
# 将文本解析为 HTML 树
|
||
# html_content = resp.text
|
||
tree = etree.HTML(html_content)
|
||
res = tree.xpath(expr)
|
||
if res:
|
||
# 获取节点文本或属性值
|
||
target_node = res[index]
|
||
value = target_node.text if hasattr(target_node, 'text') else str(target_node)
|
||
else: # 正则
|
||
res = re.findall(expr, str(target_data))
|
||
if res: value = res[index]
|
||
|
||
if value is None:
|
||
logger.warning(f"变量 [{var_name}] 未通过表达式 [{expr}] 提取到数据")
|
||
value = "not data"
|
||
|
||
self._cache[var_name] = value
|
||
logger.info(f"变量提取成功: {var_name} -> {value} (Type: {type(value).__name__})")
|
||
|
||
except Exception as e:
|
||
logger.error(f"提取变量 [{var_name}] 过程中发生异常: {e}", exc_info=True)
|
||
|
||
def _smart_replace(self, content: Any) -> Any:
|
||
"""
|
||
递归替换逻辑：
|
||
- 如果是纯变量占位符 ${token}，则返回变量在缓存中的原始类型 (int, dict, list 等)
|
||
- 如果是混合字符串或函数调用，则调用 Template 渲染为字符串
|
||
"""
|
||
if isinstance(content, dict):
|
||
return {k: self._smart_replace(v) for k, v in content.items()}
|
||
elif isinstance(content, list):
|
||
return [self._smart_replace(i) for i in content]
|
||
elif isinstance(content, str):
|
||
# A. 场景：纯变量（为了保持类型，不走 Template 渲染成字符串）
|
||
# 例子：content = "${order_id}"，如果 order_id 是 int 123，则返回 123
|
||
full_match = self.var_only_pattern.fullmatch(content)
|
||
if full_match:
|
||
var_name = full_match.group(1)
|
||
|
||
return self._cache.get(var_name, content)
|
||
|
||
# B. 场景：混合文本或函数调用
|
||
# 例子："Bearer ${token}" 或 "${gen_phone()}"
|
||
if "${" in content:
|
||
# 调用你提供的 Template 类
|
||
return Template(content).render(self._cache)
|
||
|
||
return content
|
||
|
||
def replace(self, data: T) -> T:
|
||
"""
|
||
通用的变量替换入口
|
||
支持输入 dict, list, str 或 Pydantic Model (需先 dump)
|
||
"""
|
||
if not data:
|
||
return data
|
||
|
||
logger.debug(f"开始变量替换，原始数据类型: {type(data).__name__}")
|
||
|
||
rendered_data = self._smart_replace(data)
|
||
|
||
return rendered_data
|
||
|
||
|
||
if __name__ == "__main__":
|
||
from core.models import RawSchema, HttpAction
|
||
|
||
file_handler = YamlProcessor(filepath=EXTRACT_CACHE)
|
||
variable_cache_ = file_handler.load() or {}
|
||
|
||
ex = Exchange(variable_cache_)
|
||
|
||
|
||
# --- 场景 1: 变量提取验证 ---
|
||
class MockResponse:
|
||
def __init__(self):
|
||
self.json_data = {"data": {"token": "auth_123", "user_id": 888}}
|
||
self.text = "<html><body><div id='name'>ChenWei</div></body></html>"
|
||
|
||
def json(self): return self.json_data
|
||
|
||
|
||
mock_resp = MockResponse()
|
||
print(">>> 执行提取...")
|
||
ex.extract(mock_resp, "token", "json", "$.data.token")
|
||
ex.extract(mock_resp, "u_id", "json", "$.data.user_id")
|
||
ex.extract(mock_resp, "user_name", "text", "//div[@id='name']")
|
||
|
||
# --- 场景 2: 变量替换与类型保持 ---
|
||
# 定义一个复杂的 CaseInfo
|
||
raw_case = {
|
||
"title": "测试用例",
|
||
"action": {
|
||
"method": "POST",
|
||
"url": "http://api.com/${token}", # 混合文本 -> 应转为 str
|
||
"json_body": {
|
||
"id": "${u_id}", # 纯变量 -> 应保持 int
|
||
"name": "${user_name}", # 纯变量 -> str
|
||
"config": "${existing_var}" # 初始文件变量 -> int
|
||
},
|
||
"timeout": "${existing_var}" # 字符串形式的数字 -> Pydantic 应转回 int
|
||
}
|
||
}
|
||
|
||
print("\n>>> 执行替换...")
|
||
new_case_one = ex.replace(raw_case)
|
||
print(new_case_one)
|
||
RawSchema(**new_case_one)
|
||
print(new_case_one.get("action"))
|
||
action = HttpAction(**new_case_one.get("action"))
|
||
print(action)
|
||
# # --- 校验结果 ---
|
||
print("\n--- 验证结果 ---")
|
||
print(f"URL (混合文本): {action.url} | 类型: {type(action.url)}")
|
||
print(f"ID (类型保持): {action.json_body['id']} | 类型: {type(action.json_body['id'])}")
|
||
print(f"Timeout (自动转换): {action.timeout} | 类型: {type(action.timeout)}")
|
||
# #
|
||
assert isinstance(action.json_body['id'], int)
|
||
# #
|
||
assert action.url == "http://api.com/auth_123"
|
||
assert action.timeout == 100
|
||
|
||
print("\nExchange 场景全部验证通过！")
|