开发团队有个老毛病：写测试用例的积极性不高。业务代码写完就算交差了，测试覆盖长期停留在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}
要求
使用pytest框架，不要用unittest
每个测试用例一个独立的函数，命名格式 test_{{场景描述}}
必须覆盖以下场景：
使用pytest的parametrize装饰器来组织多个测试输入
如果函数有依赖（比如调用了其他模块），使用unittest.mock进行mock
每个测试用例加上简短的中文注释说明测试意图
直接输出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干这种重复劳动正好。
“无他，惟手熟尔”！有需要的用起来！关注微信公众号「Nicholas与Pypi」获取更多Python实战！------加入知识库与更多人一起学习------https://ima.qq.com/wiki/?shareId=f2628818f0874da17b71ffa0e5e8408114e7dbad46f1745bbd1cc1365277631c
10个月宝宝每天需要喝多少奶粉？
