开发团队有个老毛病：写测试用例的积极性不高。业务代码写完就算交差了，测试覆盖长期停留在40%左右。每次代码评审的时候都说“下次补上”，下次永远不会来。

写测试确实枯燥。尤其是一些边界值和异常路径的测试，想都想不全，写起来又费时间。一个函数三五行代码，测试写起来可能二三十行，谁都不太愿意干这事。

上个月试了一个思路：用AST解析源码里所有的函数，把函数代码喂给大模型，让它生成pytest测试用例。生成的测试直接写入tests/目录，跑一下看看覆盖率。结果出乎意料，覆盖率先从40%跳到72%，调了一轮prompt之后到了85%。

关键是省事。以前给一个模块补测试要花半天，现在跑个脚本两分钟搞定，剩下的时间只需要review和微调。

方案

整个流程分四步。第一步用Python的ast模块解析源代码文件，提取出每个函数的签名、参数、返回值类型以及函数体代码。第二步把这些信息组装成一个精心设计的prompt，告诉大模型“这是一个Python函数，请为它生成pytest测试用例，包括正常路径、边界值、异常处理”。第三步解析大模型返回的代码，提取出有效的pytest测试代码。第四步把测试代码写入tests/目录下对应的文件里。

AST解析的好处是精确。不像正则匹配那样容易漏掉东西或者匹配到不该匹配的内容，AST能准确拿到函数定义的所有信息。

完整代码

pip ins tall openai pytest pytest-cov

pytest-cov用来跑覆盖率统计。

# pip install openai pytest pytest-cov# API key替换成你自己的import astimport osimport reimport textwrapfrom pathlib import Pathfrom openai import OpenAI# ============ 第一步：AST解析源码，提取函数信息 ============class FunctionInfo:"""存储一个函数的完整信息"""def __init__(self, name: str, args: list, returns: str,source_code: str, decorators: list, class_name: str = ""):self.name = nameself.args = args# 参数列表，每个元素是 {"name": "x", "type": "int", "default": "0"}self.returns = returns# 返回值类型注解self.source_code = source_codeself.decorators = decoratorsself.class_name = class_name# 如果是类方法，记录类名@propertydef full_name(self) ->str:if self.class_name:returnf"{self.class_name}.{self.name}"return self.namedef parse_python_file(file_path: str) ->list:"""    解析一个Python文件，提取所有函数信息    用ast模块，比正则匹配靠谱得多    """with open(file_path, "r", encoding="utf-8") as f:source = f.read()try:tree = ast.parse(source)except SyntaxError as e:print(f"  语法错误，跳过文件 {file_path}: {e}")return []functions = []for node in ast.walk(tree):# 处理普通函数if isinstance(node, ast.FunctionDef):func_info = _extract_function(node, source)if func_info:functions.append(func_info)# 处理类里的方法elif isinstance(node, ast.ClassDef):class_name = node.namefor item in node.body:if isinstance(item, ast.FunctionDef):# 跳过 __init__ 这种构造器，测试意义不大if item.name.startswith("__") and item.name.endswith("__"):continuefunc_info = _extract_function(item, source, class_name)if func_info:functions.append(func_info)return functionsdef_extract_function(node: ast.FunctionDef, source: str,class_name: str = "") ->FunctionInfo:"""从AST节点提取函数详细信息"""# 提取参数信息args = []for arg in node.args.args:# 跳过self和clsif arg.arg in ("self", "cls"):continuearg_info = {"name": arg.arg, "type": "", "default": ""}# 类型注解if arg.annotation:arg_info["type"] = ast.unparse(arg.annotation)args.append(arg_info)# 提取默认值defaults = node.args.defaultsif defaults:# 默认值只对应最后几个参数offset = len(args) -len(defaults)fo ri, default in enumerate(defaults):if offset+i<len(args):args[offset+i]["default"] = ast.unparse(default)# 返回值类型注解returns = ""if node.returns:returns = ast.unparse(node.returns)# 提取装饰器decorators = []for dec in node.decorator_list:decorators.append(ast.unparse(dec))# 提取函数源码（包含缩进）func_source = ast.get_source_segment(source, node)if func_source is None:# 兜底方案：用行号范围截取lines = source.split("\n")start = node.lineno-1end = node.end_linenofunc_source = "\n".join(lines[start:end])return FunctionInfo(name=node.name,args=args,returns=returns,source_code=func_source,decorators=decorators,class_name=class_name,    )# ============ 第二步：构造prompt让大模型生成测试 ============def build_test_prompt(func: FunctionInfo) ->str:"""    根据函数信息构造prompt    prompt的质量直接决定了生成的测试用例质量    """# 函数签名描述args_desc = ", ".join(f"{a['name']}: {a['type'] or 'Any'}"+ (f" = {a['default']}"ifa['default'] else "")for a in func.args    )sig = f"def {func.name}({args_desc})"if func.returns:sig += f" -> {func.returns}"prompt = f"""你是一个资深的Python测试工程师。请为以下函数生成完整的pytest测试用例。## 函数代码```python{func.source_code}

函数签名

{sig}

要求

1. 使用pytest框架，不要用unittest

2. 每个测试用例一个独立的函数，命名格式 test_{{场景描述}}

3. 必须覆盖以下场景：

• 正常输入（至少2个测试用例，使用不同的输入值）

• 边界值（空字符串、空列表、0、None、极大值等，根据参数类型选择适用的）

• 异常输入（类型错误、超出范围的值等，用pytest.raises验证）

• 如果函数有条件分支，每个分支至少一个测试

4. 使用pytest的parametrize装饰器来组织多个测试输入

5. 如果函数有依赖（比如调用了其他模块），使用unittest.mock进行mock

6. 每个测试用例加上简短的中文注释说明测试意图

7. 直接输出Python代码，不需要额外解释

输出格式

import pytest# ... 其他import# 测试代码

“””    return prompt

============ 第三步：调用大模型并解析输出 ============

class TestGenerator:“”“调用大模型生成测试用例”“”

def __init__(self):    self.client = OpenAI(        api_key="YOUR_API_KEY",        base_url="https://api.deepseek.com/v1",    )    self.model = "deepseek-chat"def generate(self, func: FunctionInfo) -> str:    """为一个函数生成测试代码"""    prompt = build_test_prompt(func)    response = self.client.chat.completions.create(        model=self.model,        messages=[            {"role": "system", "content": "你是一个专业的Python测试工程师，只输出代码，不要任何多余的解释。"},            {"role": "user", "content": prompt},        ],        temperature=0.3,  # 温度低一点，生成的代码更稳定        max_tokens=4096,    )    raw_reply = response.choices[0].message.content    # 从回复里提取代码块    code = self._extract_code(raw_reply)    return codedef _extract_code(self, text: str) -> str:    """    从大模型的回复中提取Python代码    有时候回复会带markdown代码块标记，需要去掉    """    # 尝试匹配 ```python ... ``` 代码块    pattern = r"```python\s*\n(.*?)```"    matches = re.findall(pattern, text, re.DOTALL)    if matches:        # 如果有多个代码块，全部拼起来        return "\n\n".join(matches)    # 如果没有代码块标记，看看整体是不是就是代码    # 简单判断：如果第一行是import或者def，就认为是纯代码    stripped = text.strip()    if stripped.startswith(("import", "from", "def", "#", "@pytest")):        return stripped    # 实在提取不出来就原样返回    return stripped

============ 第四步：写入测试文件 ============

def write_test_file(source_file: str, functions: list, test_codes: list,                    output_dir: str = “tests”):“””    把生成的测试代码写入文件    文件命名和源文件对应：src/utils.py -> tests/test_utils.py“””    os.makedirs(output_dir, exist_ok=True)

source_name = Path(source_file).stemtest_file = os.path.join(output_dir, f"test_{source_name}.py")lines = []lines.append(f'"""自动生成的测试用例 - 源文件: {source_file}"""')lines.append("")# 收集所有测试代码里用到的importall_imports = set()test_bodies = []for code in test_codes:    if not code.strip():        continue    # 把import行和非import行分开    import_lines = []    body_lines = []    for line in code.split("\n"):        stripped = line.strip()        if stripped.startswith(("import ", "from ")):            all_imports.add(stripped)        else:            body_lines.append(line)    test_bodies.append("\n".join(body_lines))# 先写import# 固定要加的几个importall_imports.add("import pytest")for imp in sorted(all_imports):    lines.append(imp)lines.append("")lines.append("")# 再写测试函数for body in test_bodies:    lines.append(body)    lines.append("")    lines.append("")content = "\n".join(lines)with open(test_file, "w", encoding="utf-8") as f:    f.write(content)print(f"  测试文件已写入: {test_file} ({len(test_codes)} 个函数的测试)")return test_file

============ 第五步：跑测试和覆盖率 ============

def run_tests_with_coverage(test_dir: str = “tests”, source_dir: str = “.”):“”“跑pytest并统计覆盖率”“”    import subprocess

cmd = [    "python", "-m", "pytest",    test_dir,    f"--cov={source_dir}",    "--cov-report=term-missing",    "-v",    "--tb=short",]print(f"\n运行测试: {' '.join(cmd)}\n")result = subprocess.run(cmd, capture_output=True, text=True)print(result.stdout)if result.stderr:    print("警告/错误:", result.stderr)return result.returncode

============ 主流程 ============

def generate_tests_for_project(source_dir: str, output_dir: str = “tests”):“””    扫描项目目录，为所有Python文件生成测试    这是主入口函数“””    generator = TestGenerator()

# 收集所有Python文件py_files = []for root, dirs, files in os.walk(source_dir):    # 跳过测试目录和虚拟环境    dirs[:] = [d for d in dirs if d not in ("tests", "test", "__pycache__", "venv", ".venv", "node_modules")]    for f in files:        if f.endswith(".py") and not f.startswith("test_"):            py_files.append(os.path.join(root, f))print(f"找到 {len(py_files)} 个Python源文件\n")total_functions = 0total_tests_generated = 0for file_path in py_files:    print(f"处理: {file_path}")    # 解析函数    functions = parse_python_file(file_path)    if not functions:        print("  没有可测试的函数，跳过")        continue    print(f"  发现 {len(functions)} 个函数")    total_functions += len(functions)    # 为每个函数生成测试    test_codes = []    for func in functions:        print(f"  生成测试: {func.full_name}...", end=" ")        try:            code = generator.generate(func)            test_codes.append(code)            print("完成")        except Exception as e:            print(f"失败: {e}")            test_codes.append("")    # 写入文件    valid_codes = [c for c in test_codes if c.strip()]    if valid_codes:        write_test_file(file_path, functions, test_codes, output_dir)        total_tests_generated += len(valid_codes)print(f"\n生成完成：")print(f"  处理了 {len(py_files)} 个文件")print(f"  提取了 {total_functions} 个函数")print(f"  成功生成 {total_tests_generated} 个函数的测试用例")

if name == “main”:    import sys

source = sys.argv[1] if len(sys.argv) > 1 else "."output = sys.argv[2] if len(sys.argv) > 2 else "tests"print(f"源文件目录: {source}")print(f"测试输出目录: {output}")print()# 第一步：生成测试generate_tests_for_project(source, output)# 第二步：跑测试看覆盖率print("\n" + "=" * 50)print("接下来跑测试看覆盖率...")print("=" * 50)run_tests_with_coverage(output, source)

## 运行效果找一个你自己的项目目录跑一下：```bashpython generate_tests.py ./my_project/ ./my_project/tests/

输出大概是这样：

找到 8 个Python源文件处理: ./my_project/utils.py  发现 5 个函数  生成测试: calculate_discount... 完成  生成测试: parse_user_input... 完成  生成测试: retry_with_backoff... 完成  生成测试: merge_dicts... 完成  生成测试: validate_email... 完成  测试文件已写入: tests/test_utils.py (5 个函数的测试)处理: ./my_project/services/user_service.py  发现 3 个函数  生成测试: create_user... 完成  生成测试: get_user_by_id... 完成  生成测试: update_user_profile... 完成  测试文件已写入: tests/test_user_service.py (3 个函数的测试)...生成完成：  处理了 8 个文件  提取了 23 个函数  成功生成 21 个函数的测试用例==================================================接下来跑测试看覆盖率...==================================================

pytest跑完之后会输出覆盖率报告：

tests/test_utils.py::test_calculate_discount_normal PASSEDtests/test_utils.py::test_calculate_discount_zero_price PASSEDtests/test_utils.py::test_calculate_discount_negative PASSEDtests/test_utils.py::test_parse_user_input_valid PASSEDtests/test_utils.py::test_parse_user_input_empty PASSED...Name                        Stmts   Miss  Cover   Missingmy_project/utils.py            45      3    93%   78-80my_project/services/user_service.py  62     12    81%   45-48, 92-95...TOTAL                         380     57    85%

从我自己项目的测试数据来看，之前手动写的测试覆盖率大概在40%左右。AI生成的测试加上去之后到了85%。剩下的15%主要是一些复杂的集成场景和需要真实数据库连接的测试，这些AI搞不定，还是得手动写。

生成的测试长什么样

拿一个简单的函数举例。源码是这样的：

def calculate_discount(price: float, discount_rate: float) ->float:if discount_rate<0 or discount_rate>1:raise ValueError("折扣率必须在0到1之间")if price<0:raise ValueError("价格不能为负数")return price* (1-discount_rate)

AI生成的测试是这样的：

import pytestfrom utils import calculate_discount# 正常折扣计算@pytest.mark.parametrize("price,rate,expected", [    (100.0, 0.1, 90.0),     # 打九折    (200.0, 0.5, 100.0),    # 打五折    (50.0, 0.0, 50.0),      # 不打折    (100.0, 1.0, 0.0),      # 免费])def test_calculate_discount_normal(price, rate, expected):"""正常折扣场景，验证计算结果"""assert calculate_discount(price, rate) == expected# 边界值：价格为0def test_calculate_discount_zero_price():"""价格为0时应该返回0"""assert calculate_discount(0, 0.5) == 0.0# 异常：折扣率超出范围@pytest.mark.parametrize("rate", [-0.1, 1.1, 2.0, -1.0])def test_calculate_discount_invalid_rate(rate):"""折扣率不在0-1范围内应该抛异常"""with pytest.raises(ValueError, match="折扣率必须在0到1之间"):calculate_discount(100.0, rate)# 异常：价格为负数def test_calculate_discount_negative_price():"""价格不能为负数"""with pytest.raises(ValueError, match="价格不能为负数"):calculate_discount(-10.0, 0.5)# 异常：传入非数字类型@pytest.mark.parametrize("price,rate", [    ("100", 0.5),    (None, 0.1),    (100, "0.5"),])def test_calculate_discount_type_error(price, rate):"""非数字类型应该抛出TypeError"""with pytest.raises(TypeError):calculate_discount(price, rate)

测试分成了正常路径、边界值、异常输入三个维度，用parametrize来组织多组输入。说实话，比我手写的还全。我自己写测试经常忘掉“价格为0”或者“传入None”这种情况。

实际使用中的注意事项

生成的测试不一定都能直接跑。有些函数依赖外部服务或者数据库，AI没法完全了解你的运行环境，生成的mock可能不太对。我的做法是先生成，然后花十几分钟review一遍，把有问题的mock修一下，删掉明显不合理的用例。

prompt里的温度设成0.3比较合适。温度太高生成的代码花样多但容易出语法错误，太低的话测试用例容易重复。0.3到0.5之间效果最好。

还有一个技巧：如果你的项目有类型注解，AST提取出来的信息更丰富，大模型生成的测试质量也更高。没有类型注解的函数，大模型只能猜参数是什么类型，有时候会猜错。

小结

AST解析加大模型生成，几分钟就能把一个项目的测试覆盖率从40%拉到85%。写测试这事终于不用再靠毅力了，让AI干这种重复劳动正好。

https://ima.qq.com/wiki/?shareId=f2628818f0874da17b71ffa0e5e8408114e7dbad46f1745bbd1cc1365277631c