如何从0构建一款类似pytest的工具

Pytest主要模块

从0构建一个类似pytest的工具

  前面简要介绍了pytest的主要功能模块，如果要从0构建一个类似pytest的工具，应该如何实现呢？下面是实现的具体代码。

import os
import importlib.util
import inspect
import traceback
import argparse

# 发现测试文件
def discover_tests(start_dir):
    test_files = []
    for root, _, files in os.walk(start_dir):
        for file in files:
            if file.startswith('test_') and file.endswith('.py'):
                test_files.append(os.path.join(root, file))
    return test_files

# 查找测试函数
def find_test_functions(module):
    test_functions = []
    for name, obj in inspect.getmembers(module):
        if inspect.isfunction(obj) and name.startswith('test_'):
            test_functions.append(obj)
    return test_functions

# 运行测试函数
def run_tests(test_functions):
    results = []
    for test_func in test_functions:
        result = {'name': test_func.__name__}
        try:
            test_func()
            result['status'] = 'pass'
        except AssertionError as e:
            result['status'] = 'fail'
            result['error'] = traceback.format_exc()
        except Exception as e:
            result['status'] = 'error'
            result['error'] = traceback.format_exc()
        results.append(result)
    return results

# 打印测试结果
def print_results(results):
    for result in results:
        print(f'Test: {result["name"]} - {result["status"]}')
        if result.get('error'):
            print(result['error'])
        print('-' * 40)

# 主函数
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='A simple pytest-like tool')
    parser.add_argument('test_path', type=str, help='Path to the test file or directory')

    args = parser.parse_args()
    test_path = args.test_path

    if os.path.isdir(test_path):
        test_files = discover_tests(test_path)
    elif os.path.isfile(test_path):
        test_files = [test_path]
    else:
        print(f"Invalid path: {test_path}")
        exit(1)

    for test_file in test_files:
    # 根据测试文件路径创建模块规范
      spec = importlib.util.spec_from_file_location("module.name", test_file)
    
    # 根据模块规范创建一个模块对象
      module = importlib.util.module_from_spec(spec)
    
    # 加载并执行模块代码
      spec.loader.exec_module(module)
    
    # 在模块中查找测试函数
      test_functions = find_test_functions(module)
    
    # 运行所有找到的测试函数，并记录结果
      results = run_tests(test_functions)
    
    # 输出测试结果
      print_results(results)

  准备测试脚本文件:test_example.py，内容如下所示：每个测试方法都是以test开头，这样上面的代码才能正确捕获到测试方法。

def test_addition():
    assert 1 + 1 == 2

def test_subtraction():
    assert 2 - 1 == 1

def test_failure():
    assert 1 + 1 == 3

  执行命令"python3 simple_pytest.py test_example.py"，运行测试，结果如下：两个执行成功，一个失败。说明整个工具功能符合预期。

 importlib.util包

上面的代码通过importlib.util来动态加载和操作模块，importlib.util的主要作用是提供实用工具来帮助开发者在运行时动态加载模块，而不是在编译时静态加载。这对于需要在程序执行期间动态加载模块的场景非常有用，例如插件系统、测试框架等。提供的主要方法有：

  spec_from_file_location(name, location, *, loader=None, submodule_search_locations=None)： 根据文件路径创建一个模块规范 (ModuleSpec)
  module_from_spec(spec)：根据模块规范创建一个新的模块对象
  spec.loader.exec_module(module)：执行加载模块的代码，将代码注入到模块对象中
  find_spec(name, package=None)：查找指定名称的模块规范

  模块规范具体包含哪些属性呢？模块规范主要包含模块名称，模块的加载器，模块的来源，是否有文件路径，子模块搜索路径，缓存路径，是否有父模块。

  下面这段代码演示了如何通过importlib.util包来创建模块，并调用模块中的函数。

import importlib.util
# 获取模块文件路径
file_path = "example_module.py"
# 创建模块规范对象
spec = importlib.util.spec_from_file_location("example_module", file_path)
# 打印ModuleSpec对象的信息
print("ModuleSpec Information:")
print(f"Name: {spec.name}")
print(f"Loader: {spec.loader}")
print(f"Origin: {spec.origin}")
print(f"Has Location: {spec.has_location}")
print(f"Submodule Search Locations: {spec.submodule_search_locations}")
# 创建模块对象
module = importlib.util.module_from_spec(spec)
# 加载并执行模块
spec.loader.exec_module(module)
# 调用模块中的函数
module.hello()
module.test_addition()
module.test_failure()

example_module.py测试文件内容

def hello():
    print("Hello from example_module!")

def test_addition():
    assert 1 + 1 == 2 

def test_failure():
    assert 1 + 1 == 3    

  执行结果如下所示：可以看到文件中的函数都被执行了，且给出了执行结果。如果是测试框架，就可以收集这些测试结果，用户后续的测试报告显示。

自定义命令运行测试文件

   前面在执行测试的时候，是通过python命令来执行测试文件的，如果要像pytest一样，通过自定义命令来执行测试文件，应该如何实现呢？这里需要借助Python的setuptools包中的 entry_points 功能。通过定义一个控制台脚本，让用户直接通过命令行运行工具。在原来代码基础上，创建setup.py文件。entry_points中console_scripts中，定义了自定义命令是my_pytests,对应的代码入口是之前的工具实现文件simple_pytest文件中main方法。

from setuptools import setup, find_packages

setup(
    name='my_pytest',
    version='0.1',
    packages=find_packages(),
    entry_points={
        'console_scripts': [
            'my_pytests=simple_pytest:main',
        ],
    },
    python_requires='>=3.6',
)

  定义好setup文件后，通过命令进行打包"pip install -e .",就可以通过my_pytests命令执行文件了，例如“my_pytests ./test_example.py” or "my_pytests ./tests".执行结果如下所示：

  以上就是构建类似pytest工具的实现过程以及原理。