|
| 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 | +
|
| 35 | + # 在模块导入时就初始化，确保在 httpx 被导入之前完成 patch |
| 36 | + initialize_header_injection() |
| 37 | +
|
| 38 | + except ImportError: |
| 39 | + # testsolar_pytestx 未安装，跳过 |
| 40 | + pass |
| 41 | +''' |
| 42 | + |
| 43 | +# 用于标识自动生成的 conftest.py 的标记 |
| 44 | +CONFTEST_MARKER = "Auto-generated conftest.py for testsolar header injection support." |
| 45 | + |
| 46 | + |
| 47 | +def _is_testsolar_conftest(conftest_path: Path) -> bool: |
| 48 | + """检查 conftest.py 是否是由 testsolar 生成的""" |
| 49 | + if not conftest_path.exists(): |
| 50 | + return False |
| 51 | + |
| 52 | + try: |
| 53 | + content = conftest_path.read_text(encoding="utf-8") |
| 54 | + return CONFTEST_MARKER in content |
| 55 | + except Exception: |
| 56 | + return False |
| 57 | + |
| 58 | + |
| 59 | +def _backup_existing_conftest(conftest_path: Path) -> Optional[Path]: |
| 60 | + """备份已存在的 conftest.py（如果不是 testsolar 生成的）""" |
| 61 | + if not conftest_path.exists(): |
| 62 | + return None |
| 63 | + |
| 64 | + if _is_testsolar_conftest(conftest_path): |
| 65 | + # 是 testsolar 生成的，不需要备份 |
| 66 | + return None |
| 67 | + |
| 68 | + # 创建备份 |
| 69 | + backup_path = conftest_path.with_suffix(".py.testsolar_backup") |
| 70 | + try: |
| 71 | + import shutil |
| 72 | + |
| 73 | + shutil.copy2(conftest_path, backup_path) |
| 74 | + logger.info(f"Backed up existing conftest.py to {backup_path}") |
| 75 | + return backup_path |
| 76 | + except Exception as e: |
| 77 | + logger.error(f"Failed to backup conftest.py: {e}") |
| 78 | + return None |
| 79 | + |
| 80 | + |
| 81 | +def _merge_conftest_content(existing_content: str, template_content: str) -> str: |
| 82 | + """合并已有的 conftest.py 内容和模板内容""" |
| 83 | + # 如果已有内容中已经包含了我们的标记，说明已经合并过了 |
| 84 | + if CONFTEST_MARKER in existing_content: |
| 85 | + return existing_content |
| 86 | + |
| 87 | + # 将模板内容添加到已有内容的开头 |
| 88 | + merged = ( |
| 89 | + template_content + "\n\n# === Original conftest.py content below ===\n\n" + existing_content |
| 90 | + ) |
| 91 | + return merged |
| 92 | + |
| 93 | + |
| 94 | +def generate_conftest_for_header_injection(project_path: str) -> Optional[str]: |
| 95 | + """ |
| 96 | + 在项目目录中生成支持 header injection 的 conftest.py |
| 97 | +
|
| 98 | + Args: |
| 99 | + project_path: 项目根目录路径 |
| 100 | +
|
| 101 | + Returns: |
| 102 | + 生成的 conftest.py 路径，如果生成失败则返回 None |
| 103 | + """ |
| 104 | + conftest_path = Path(project_path) / "conftest.py" |
| 105 | + |
| 106 | + try: |
| 107 | + if conftest_path.exists(): |
| 108 | + if _is_testsolar_conftest(conftest_path): |
| 109 | + # 已经是 testsolar 生成的，直接覆盖更新 |
| 110 | + logger.debug(f"Updating existing testsolar conftest.py at {conftest_path}") |
| 111 | + conftest_path.write_text(CONFTEST_TEMPLATE, encoding="utf-8") |
| 112 | + else: |
| 113 | + # 已有用户的 conftest.py，需要合并 |
| 114 | + logger.info(f"Found existing conftest.py at {conftest_path}, merging content") |
| 115 | + existing_content = conftest_path.read_text(encoding="utf-8") |
| 116 | + merged_content = _merge_conftest_content(existing_content, CONFTEST_TEMPLATE) |
| 117 | + |
| 118 | + # 备份原文件 |
| 119 | + _backup_existing_conftest(conftest_path) |
| 120 | + |
| 121 | + # 写入合并后的内容 |
| 122 | + conftest_path.write_text(merged_content, encoding="utf-8") |
| 123 | + logger.info("Merged testsolar header injection into existing conftest.py") |
| 124 | + else: |
| 125 | + # 不存在 conftest.py，直接创建 |
| 126 | + logger.info(f"Creating conftest.py at {conftest_path}") |
| 127 | + conftest_path.write_text(CONFTEST_TEMPLATE, encoding="utf-8") |
| 128 | + |
| 129 | + return str(conftest_path) |
| 130 | + |
| 131 | + except Exception as e: |
| 132 | + logger.error(f"Failed to generate conftest.py: {e}") |
| 133 | + return None |
| 134 | + |
| 135 | + |
| 136 | +def cleanup_generated_conftest(project_path: str) -> None: |
| 137 | + """ |
| 138 | + 清理生成的 conftest.py |
| 139 | +
|
| 140 | + 如果 conftest.py 是纯粹由 testsolar 生成的，则删除它。 |
| 141 | + 如果是合并的，则恢复原始内容。 |
| 142 | +
|
| 143 | + Args: |
| 144 | + project_path: 项目根目录路径 |
| 145 | + """ |
| 146 | + conftest_path = Path(project_path) / "conftest.py" |
| 147 | + backup_path = conftest_path.with_suffix(".py.testsolar_backup") |
| 148 | + |
| 149 | + try: |
| 150 | + if backup_path.exists(): |
| 151 | + # 存在备份，恢复原始文件 |
| 152 | + import shutil |
| 153 | + |
| 154 | + shutil.move(str(backup_path), str(conftest_path)) |
| 155 | + logger.info("Restored original conftest.py from backup") |
| 156 | + elif conftest_path.exists() and _is_testsolar_conftest(conftest_path): |
| 157 | + # 是纯粹由 testsolar 生成的，检查是否只有模板内容 |
| 158 | + content = conftest_path.read_text(encoding="utf-8") |
| 159 | + if "Original conftest.py content below" not in content: |
| 160 | + # 没有合并用户内容，可以安全删除 |
| 161 | + conftest_path.unlink() |
| 162 | + logger.info("Removed generated conftest.py") |
| 163 | + except Exception as e: |
| 164 | + logger.warning(f"Failed to cleanup conftest.py: {e}") |
0 commit comments