Skip to content

Commit d49cd3b

Browse files
ambzhangambzhang
authored
Merge branch 'main' into dev/fix-parallel-bug
2 parents 394b5c1 + b2c70e6 commit d49cd3b

8 files changed

Lines changed: 842 additions & 2 deletions

File tree

‎pytest/CHANGELOG.md‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1212
- Refactor reporter implementation to use FileReporter instead of Reporter
1313

1414
### Fixed
15-
- Improve test case collection and execution with better file handling
15+
- Improve test case collection and execution with better file handling
16+
17+
### Added
18+
- Support for HTTP request header injection during test execution
19+
- Add test case identification to API requests via X-Testsolar-Testcase header

‎pytest/src/testsolar_pytestx/conftest_generator.py‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
"""
2+
conftest.py 生成器模块
3+
4+
用于在测试目录中生成 conftest.py 文件，确保 xdist worker 进程也能正确初始化 header injection。
5+
6+
当使用 pytest.main(plugins=[PytestExecutor]) 方式运行测试时，
7+
xdist worker 进程不会继承主进程的 plugins 参数，
8+
因此需要通过 conftest.py 来确保每个 worker 进程都能执行初始化。
9+
"""
10+
11+
from pathlib import Path
12+
from typing import Optional
13+
14+
from loguru import logger
15+
16+
# conftest.py 的内容模板
17+
# 注意：这个模板会被写入到测试项目目录中
18+
CONFTEST_TEMPLATE = '''"""
19+
Auto-generated conftest.py for testsolar header injection support.
20+
This file ensures that HTTP request header injection works correctly
21+
even when using pytest-xdist for parallel test execution.
22+
23+
DO NOT MODIFY THIS FILE - it will be regenerated on each test run.
24+
"""
25+
import os
26+
27+
# 仅在启用 API 收集时才初始化
28+
if os.environ.get("ENABLE_API_COLLECTING") == "1":
29+
try:
30+
from testsolar_pytestx.header_injection import (
31+
initialize_header_injection,
32+
set_current_test_nodeid,
33+
)
34+
from testsolar_pytestx.converter import normalize_testcase_name
35+
36+
# 在模块导入时就初始化，确保在 httpx 被导入之前完成 patch
37+
initialize_header_injection()
38+
39+
def pytest_runtest_setup(item):
40+
"""在每个测试开始前设置 nodeid"""
41+
testcase_name = normalize_testcase_name(item.nodeid)
42+
testcase_class_name = testcase_name.split("?", 1)[-1]
43+
set_current_test_nodeid(testcase_class_name)
44+
45+
def pytest_runtest_teardown(item):
46+
"""在每个测试结束后清除 nodeid"""
47+
set_current_test_nodeid(None)
48+
49+
except ImportError:
50+
# testsolar_pytestx 未安装，跳过
51+
pass
52+
'''
53+
54+
# 用于标识自动生成的 conftest.py 的标记
55+
CONFTEST_MARKER = "Auto-generated conftest.py for testsolar header injection support."
56+
57+
58+
def _is_testsolar_conftest(conftest_path: Path) -> bool:
59+
"""检查 conftest.py 是否是由 testsolar 生成的"""
60+
if not conftest_path.exists():
61+
return False
62+
63+
try:
64+
content = conftest_path.read_text(encoding="utf-8")
65+
return CONFTEST_MARKER in content
66+
except Exception:
67+
return False
68+
69+
70+
def _backup_existing_conftest(conftest_path: Path) -> Optional[Path]:
71+
"""备份已存在的 conftest.py（如果不是 testsolar 生成的）"""
72+
if not conftest_path.exists():
73+
return None
74+
75+
if _is_testsolar_conftest(conftest_path):
76+
# 是 testsolar 生成的，不需要备份
77+
return None
78+
79+
# 创建备份
80+
backup_path = conftest_path.with_suffix(".py.testsolar_backup")
81+
try:
82+
import shutil
83+
84+
shutil.copy2(conftest_path, backup_path)
85+
logger.info(f"Backed up existing conftest.py to {backup_path}")
86+
return backup_path
87+
except Exception as e:
88+
logger.error(f"Failed to backup conftest.py: {e}")
89+
return None
90+
91+
92+
def _merge_conftest_content(existing_content: str, template_content: str) -> str:
93+
"""合并已有的 conftest.py 内容和模板内容"""
94+
# 如果已有内容中已经包含了我们的标记，说明已经合并过了
95+
if CONFTEST_MARKER in existing_content:
96+
return existing_content
97+
98+
# 将模板内容添加到已有内容的开头
99+
merged = (
100+
template_content + "\n\n# === Original conftest.py content below ===\n\n" + existing_content
101+
)
102+
return merged
103+
104+
105+
def generate_conftest_for_header_injection(project_path: str) -> Optional[str]:
106+
"""
107+
在项目目录中生成支持 header injection 的 conftest.py
108+
109+
Args:
110+
project_path: 项目根目录路径
111+
112+
Returns:
113+
生成的 conftest.py 路径，如果生成失败则返回 None
114+
"""
115+
conftest_path = Path(project_path) / "conftest.py"
116+
117+
try:
118+
if conftest_path.exists():
119+
if _is_testsolar_conftest(conftest_path):
120+
# 已经是 testsolar 生成的，直接覆盖更新
121+
logger.debug(f"Updating existing testsolar conftest.py at {conftest_path}")
122+
conftest_path.write_text(CONFTEST_TEMPLATE, encoding="utf-8")
123+
else:
124+
# 已有用户的 conftest.py，需要合并
125+
logger.info(f"Found existing conftest.py at {conftest_path}, merging content")
126+
existing_content = conftest_path.read_text(encoding="utf-8")
127+
merged_content = _merge_conftest_content(existing_content, CONFTEST_TEMPLATE)
128+
129+
# 备份原文件
130+
_backup_existing_conftest(conftest_path)
131+
132+
# 写入合并后的内容
133+
conftest_path.write_text(merged_content, encoding="utf-8")
134+
logger.info("Merged testsolar header injection into existing conftest.py")
135+
else:
136+
# 不存在 conftest.py，直接创建
137+
logger.info(f"Creating conftest.py at {conftest_path}")
138+
conftest_path.write_text(CONFTEST_TEMPLATE, encoding="utf-8")
139+
140+
return str(conftest_path)
141+
142+
except Exception as e:
143+
logger.error(f"Failed to generate conftest.py: {e}")
144+
return None
145+
146+
147+
def cleanup_generated_conftest(project_path: str) -> None:
148+
"""
149+
清理生成的 conftest.py
150+
151+
如果 conftest.py 是纯粹由 testsolar 生成的，则删除它。
152+
如果是合并的，则恢复原始内容。
153+
154+
Args:
155+
project_path: 项目根目录路径
156+
"""
157+
conftest_path = Path(project_path) / "conftest.py"
158+
backup_path = conftest_path.with_suffix(".py.testsolar_backup")
159+
160+
try:
161+
if backup_path.exists():
162+
# 存在备份，恢复原始文件
163+
import shutil
164+
165+
shutil.move(str(backup_path), str(conftest_path))
166+
logger.info("Restored original conftest.py from backup")
167+
elif conftest_path.exists() and _is_testsolar_conftest(conftest_path):
168+
# 是纯粹由 testsolar 生成的，检查是否只有模板内容
169+
content = conftest_path.read_text(encoding="utf-8")
170+
if "Original conftest.py content below" not in content:
171+
# 没有合并用户内容，可以安全删除
172+
conftest_path.unlink()
173+
logger.info("Removed generated conftest.py")
174+
except Exception as e:
175+
logger.warning(f"Failed to cleanup conftest.py: {e}")

‎pytest/src/testsolar_pytestx/executor.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,20 @@
3838
from .parser import parse_case_attributes
3939
from .stream import pytest_main_with_output
4040

41+
# from .header_injection import set_current_test_nodeid
42+
from .conftest_generator import generate_conftest_for_header_injection
43+
4144

4245
class RunMode(Enum):
4346
SINGLE = "single"
4447
BATCH = "batch"
4548

4649

50+
def should_enable_header_injection() -> bool:
51+
"""判断是否需要启用请求头注入功能"""
52+
return os.environ.get("ENABLE_API_COLLECTING", "") == "1"
53+
54+
4755
class PytestExecutor:
4856
def __init__(
4957
self,
@@ -67,6 +75,13 @@ def pytest_runtest_logstart(self, nodeid: str, location: Any) -> None:
6775
# 通知ResultHouse用例开始运行
6876
testcase_name = normalize_testcase_name(nodeid, self.data_drive_key)
6977

78+
# 设置当前测试用例的nodeid到上下文
79+
# 注意：这里使用处理后的 testcase_class_name，与 conftest.py 中的格式可能不同
80+
# conftest.py 中使用原始 nodeid，这里使用处理后的名称
81+
# if should_enable_header_injection():
82+
# testcase_class_name = testcase_name.split("?", 1)[-1]
83+
# set_current_test_nodeid(testcase_class_name)
84+
7085
test_result = TestResult(
7186
Test=TestCase(Name=testcase_name),
7287
ResultType=ResultType.RUNNING,
@@ -210,6 +225,11 @@ def pytest_runtest_logfinish(self, nodeid: str, location: Any) -> None:
210225

211226
# 上报完成后测试记录就没有用了，删除以节省内存
212227
self.testdata.pop(testcase_name, None)
228+
229+
# 清除当前测试用例的nodeid
230+
# if should_enable_header_injection():
231+
# set_current_test_nodeid(None)
232+
213233
logger.info(f"E {nodeid} runtest_logfinish")
214234

215235
def pytest_sessionfinish(self, session: Session, exitstatus: int) -> None:
@@ -240,6 +260,13 @@ def run_testcases(
240260
if entry.ProjectPath not in sys.path:
241261
sys.path.insert(0, entry.ProjectPath)
242262

263+
# 如果启用了 header injection，生成 conftest.py 以支持 xdist
264+
# 这确保了 xdist worker 进程也能正确初始化 header injection
265+
if should_enable_header_injection():
266+
conftest_path = generate_conftest_for_header_injection(entry.ProjectPath)
267+
if conftest_path:
268+
logger.info(f"Generated conftest.py for header injection at {conftest_path}")
269+
243270
valid_selectors, _ = filter_invalid_selector_path(
244271
workspace=entry.ProjectPath,
245272
selectors=entry.TestSelectors,

0 commit comments

Comments
 (0)