支持请求头注入

Author: zixingdeng

pytest/CHANGELOG.md

@@ -12,4 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Refactor reporter implementation to use FileReporter instead of Reporter
 
 ### Fixed
-- Improve test case collection and execution with better file handling
+- Improve test case collection and execution with better file handling
+
+### Added
+- Support for HTTP request header injection during test execution
+- Add test case identification to API requests via X-Testsolar-Testcase header

pytest/src/testsolar_pytestx/executor.py

@@ -37,6 +37,10 @@
 from .filter import filter_invalid_selector_path
 from .parser import parse_case_attributes
 from .stream import pytest_main_with_output
+from .header_injection import (
+    initialize_header_injection,
+    set_current_test_nodeid,
+)
 
 
 class RunMode(Enum):
@@ -58,12 +62,23 @@ def __init__(
         self.comment_fields = comment_fields
         self.data_drive_key = data_drive_key
 
+        # 初始化API请求头注入功能，传入判定方法
+        initialize_header_injection(self._should_enable_header_injection)
+
+    def _should_enable_header_injection(self) -> bool:
+        """判断是否需要启用请求头注入功能"""
+        return os.environ.get("ENABLE_API_COLLECTING", "") == "1"
+
     def pytest_runtest_logstart(self, nodeid: str, location: Any) -> None:
         """
         Called at the start of running the runtest protocol for a single item.
         """
         logger.info(f"{nodeid} start")
 
+        # 设置当前测试用例的nodeid到上下文
+        if self._should_enable_header_injection():
+            set_current_test_nodeid(nodeid)
+
         # 通知ResultHouse用例开始运行
         testcase_name = normalize_testcase_name(nodeid, self.data_drive_key)
 
@@ -210,6 +225,11 @@ def pytest_runtest_logfinish(self, nodeid: str, location: Any) -> None:
 
             # 上报完成后测试记录就没有用了，删除以节省内存
             self.testdata.pop(testcase_name, None)
+
+        # 清除当前测试用例的nodeid
+        if self._should_enable_header_injection():
+            set_current_test_nodeid(None)
+
         logger.info(f"E {nodeid} runtest_logfinish")
 
     def pytest_sessionfinish(self, session: Session, exitstatus: int) -> None:

pytest/src/testsolar_pytestx/header_injection.py

@@ -0,0 +1,157 @@
+"""
+HTTP请求头注入模块
+
+该模块用于在pytest测试执行过程中，自动向HTTP请求注入测试用例标识信息。
+支持requests和httplib/http.client两种HTTP库。
+"""
+
+import threading
+from typing import Optional, Callable, Any, Dict
+
+from loguru import logger
+
+
+# ============ 上下文管理 ============
+# 用于存储当前测试用例的nodeid
+_test_context = threading.local()
+
+# 原始请求方法引用
+_original_requests_request: Optional[Callable[..., Any]] = None
+_original_httplib_request: Optional[Callable[..., Any]] = None
+
+# 用于存储判定函数的引用
+_should_enable_injection_func: Optional[Callable[[], bool]] = None
+
+
+def get_current_test_nodeid() -> Optional[str]:
+    """获取当前测试用例的nodeid"""
+    return getattr(_test_context, "nodeid", None)
+
+
+def set_current_test_nodeid(nodeid: Optional[str]) -> None:
+    """设置当前测试用例的nodeid"""
+    _test_context.nodeid = nodeid
+
+
+# ============ Header 注入逻辑 ============
+def _inject_header_to_dict(headers: Optional[Dict[str, Any]], nodeid: str) -> Dict[str, Any]:
+    """
+    向请求头字典注入X-Testsolar-Testcase
+
+    Args:
+        headers: 原始请求头字典
+        nodeid: 测试用例的nodeid
+
+    Returns:
+        注入后的请求头字典
+    """
+    if headers is None:
+        headers = {}
+    elif not isinstance(headers, dict):
+        # 如果headers不是dict，创建新dict
+        headers = dict(headers)
+    else:
+        # 创建副本避免修改原始dict
+        headers = headers.copy()
+
+    headers["X-Testsolar-Testcase"] = nodeid
+    return headers
+
+
+# ============ Monkey Patch 实现 ============
+def _patch_requests() -> None:
+    """Patch requests库的Session.request方法"""
+    try:
+        import requests  # type: ignore[import]
+
+        global _original_requests_request
+
+        if _original_requests_request is not None:
+            # 已经patch过了
+            return
+
+        _original_requests_request = requests.Session.request
+
+        def _patched_request(self: Any, method: str, url: str, **kwargs: Any) -> Any:
+            # 通过判定函数检查是否需要注入
+            if _should_enable_injection_func and _should_enable_injection_func():
+                nodeid = get_current_test_nodeid()
+                if nodeid:
+                    headers = kwargs.get("headers")
+                    kwargs["headers"] = _inject_header_to_dict(headers, nodeid)
+
+            return _original_requests_request(self, method, url, **kwargs)
+
+        requests.Session.request = _patched_request
+        logger.info("Successfully patched requests.Session.request")
+
+    except ImportError:
+        logger.debug("requests library not installed, skip patching")
+    except Exception as e:
+        logger.error(f"Failed to patch requests: {e}")
+
+
+def _patch_httplib() -> None:
+    """Patch httplib/http.client的HTTPConnection.request方法"""
+    try:
+        # Python 2/3 兼容
+        try:
+            import http.client as httplib
+        except ImportError:
+            import httplib  # type: ignore
+
+        global _original_httplib_request
+
+        if _original_httplib_request is not None:
+            # 已经patch过了
+            return
+
+        _original_httplib_request = httplib.HTTPConnection.request
+
+        def _patched_httplib_request(
+            self: Any,
+            method: str,
+            url: str,
+            body: Any = None,
+            headers: Any = None,
+            **kwargs: Any,
+        ) -> Any:
+            # 通过判定函数检查是否需要注入
+            final_headers: Any = headers
+            if _should_enable_injection_func and _should_enable_injection_func():
+                nodeid = get_current_test_nodeid()
+                if nodeid:
+                    # httplib的headers参数默认是{}而不是None
+                    if final_headers is None:
+                        final_headers = {"X-Testsolar-Testcase": nodeid}
+                    else:
+                        # 创建副本避免修改原始headers
+                        new_headers = dict(final_headers)
+                        new_headers["X-Testsolar-Testcase"] = nodeid
+                        final_headers = new_headers
+
+            return _original_httplib_request(self, method, url, body, final_headers, **kwargs)
+
+        httplib.HTTPConnection.request = _patched_httplib_request  # type: ignore[method-assign]
+        logger.info("Successfully patched httplib.HTTPConnection.request")
+
+    except ImportError:
+        logger.debug("httplib library not available, skip patching")
+    except Exception as e:
+        logger.error(f"Failed to patch httplib: {e}")
+
+
+# ============ 公共接口 ============
+def initialize_header_injection(should_enable_func: Callable[[], bool]) -> None:
+    """
+    初始化HTTP请求头注入功能
+
+    Args:
+        should_enable_func: 判定函数，返回bool表示是否启用注入
+    """
+    global _should_enable_injection_func
+    _should_enable_injection_func = should_enable_func
+
+    logger.info("Initializing API header injection...")
+    _patch_requests()
+    _patch_httplib()

pytest/tests/test_run.py

@@ -1,4 +1,5 @@
 import json
+import os
 import tempfile
 from pathlib import Path
 from unittest import TestCase
@@ -15,6 +16,10 @@
 class TestExecuteEntry(TestCase):
     testdata_dir = Path(__file__).parent.parent.absolute().joinpath("testdata")
 
+    def test_header_injection(self):
+        os.environ["ENABLE_API_COLLECTING"] = "1"
+        self.test_run_testcases_from_args()
+
     def test_run_testcases_from_args(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             report_dir = Path(tmpdir)