Pytest黑魔法之动态重载test case

前言

首先介绍一下背景。我们组的服务需要用python脚本来做业务逻辑相对复杂的integration tests,为此维护了一个代码库来管理这些测试代码。这套测试重度耦合了亚麻内部的一个极度复杂的pytest框架,该框架起到的作用是创建和清理测试运行过程中所需要用到的AWS资源(比如EC2 instance、VPC等等),但对于我们的使用场景来说非但没有解决什么问题,反倒是引入了一堆麻烦,但由于是为了支持做这个框架的组的KPI,属于内部政治需求不得不硬着头皮也要用。

这个内部框架严重拖累了我们这些GA重压之下的苦逼搬砖人:由于它负责管理资源,所以每次跑测试都要执行一系列复杂的步骤之后才会真的开始运行测试代码,这个过程可以从几十秒到十几分钟不等,取决于运行测试的AWS账号里的资源状态是否已经ready。让体验更加恶化的是,亚麻的代码库统一使用一个名为brazil的java实现的构建系统,这玩意构建python的时候也会执行一些标准的打包artifacts的操作,从而生成一个可以直接部署上线的runtime package,所以这个过程中会把所有的py源文件都拷贝到新的目录下,那么如果要手动执行代码库里某些测试脚本,brazil会从打包之后的环境里进行调用,而不是直接执行repo里的python脚本。

我们都知道用python这种动态类型语言写测试,代码虽然非常简短,但一大痛点在于:仅凭简短的context你很难搞明白一个对象到底是什么类型,具有哪些属性和方法(对于AWS的python client这个问题尤其严重),所以经常容易需要修修改改。这几个因素综合起来,导致编写integration test要忍受极其低效的feedback loop:对一行测试代码的改动需要两分钟的build再加上几分钟到十几分钟不等的测试框架初始化,可谓是糟心到极点。于是,作为信奉效率就是生命的程序猿,不由得开始思考——有没有什么办法能在test case不符合预期的时候把程序停下来,手动修改并重新加载代码然后从这里再次运行呢?

从技术角度,得益于python动态语言的特性和pytest的插件机制,我们完全可以在test case运行不符合预期的时候把程序停在这里,手动进行修改调整之后,直接reload失败的测试代码然后重跑这个case;而只要正确修改sys.path里的搜索路径,就可以让python runtime忽略build artifacts里面的源文件,直接加载我们修改过后的代码。这两个功能只需要通过写一个pytest插件就能实现,从而大幅度提高开发效率。

fix import path

首先我们需要更新import path,大概原理就是取出当前正在运行的脚本的绝对路径,也就是build artifacts所在的目录,然后按照约定转换成代码repo的实际路径即可,这样当下次我们试图import module的时候python会读取我们更新过的test case:

def _redirect_import_search_path_to_brazil_source():
    global _path_redirected
    if _path_redirected:
        return

    def _redirect_path(path):
        """
        redirect build artifacts path like `/local/home/<username>/workplace/MercuryControlPlaneIntegrationTestsEnv/build/MercuryControlPlaneIntegrationTests/MercuryControlPlaneIntegrationTests-1.0/AL2_x86_64/DEV.STD.PTHREAD/build/test-integ/MercuryAPI`
        to its corresponding source dir for module import.
        """
        basedir, file = os.path.split(os.path.abspath(__file__))
        repo_root = os.path.abspath(os.path.join(basedir, ".."))
        repo_name = os.path.basename(repo_root)
        if repo_name in path:
            path_elements = path.split(os.sep)
            for i in range(len(path_elements) - 1, 0, -1):
                new_path = os.path.join(repo_root, os.sep.join(path_elements[i:]))
                if os.path.exists(new_path) and os.path.isdir(new_path) and new_path != path:
                    return new_path
            return path
        else:
            return path

    sys.path[:] = [
        _redirect_path(p) for p in sys.path
    ]
    _path_redirected = True

Pytest Hook

From ChatGPT:

pytest中,hook机制是一种插件系统,它允许用户通过自定义钩子函数来扩展和修改pytest的行为。pytest会在测试运行的不同阶段触发这些钩子函数,你可以通过实现或覆盖这些钩子函数来插入自定义逻辑。

常见钩子函数:

  • pytest_configure(config):初始化配置时调用。
  • pytest_collection_modifyitems(session, config, items):修改收集到的测试用例。
  • pytest_runtest_setup(item):在每个测试用例运行前调用。
  • pytest_runtest_call(item):运行测试用例时调用。
  • pytest_runtest_teardown(item):在每个测试用例运行后调用。

pytest的钩子函数中,item 是一个 pytest.Item 对象,代表一个测试用例。它封装了关于该测试用例的各种信息,比如名称、执行函数、路径、所属模块、调用参数等。你可以通过访问 item 的属性或方法获取或操作该测试用例的相关信息。

pytest.Item 对象中,functionobj 都与测试用例的 Python 函数有关,但有细微区别:

  • function:表示包装后的测试函数对象,通常用于执行测试,它可能包含 pytest 的增强功能,如 pytest.mark 处理等。
  • obj:表示原始的测试函数对象,即未经过 pytest 处理的 Python 函数。

本文我们用到的是pytest_runtest_protocol(item, nextitem)这个钩子,它用于控制单个测试用例的执行流程,提供了一个非常灵活的机制,允许用户在测试用例执行的各个阶段插入自定义逻辑,通常用于修改测试的执行顺序、条件判断或其他特定需求。为了适应我们的使用场景,可以在这个钩子里修改执行逻辑,当test case failed时询问用户是否要修改代码然后重试即可。

test code reload

接下来我们需要在修改了代码之后重新加载它,下面这部分是reload test case的核心逻辑,整个流程trick还是比较多的:

def load_method_from_classes(module, fname):
    for name, obj in inspect.getmembers(module):
        if inspect.isclass(obj) and hasattr(obj, fname):
            return getattr(obj, fname)
    raise RuntimeError("Error: [{}] not found, did you delete it?".format(fname))

def _reload_test_code(item):
    # 首先把import path重定位到源码repo里面去,而不是build artifacts
    _redirect_import_search_path_to_brazil_source()
    # 拿到当前执行失败的test case所在的module,称之为old module
    old_module = sys.modules[item.function.__module__]

    globals_backup = {}
    # 需要备份old module下面所有的全局变量,测试执行时要用,但是要注意忽略掉那些特殊命名的
    for name, obj in inspect.getmembers(old_module):
        if not inspect.isclass(obj) and \
            not inspect.isfunction(obj) and \
            not inspect.ismodule(obj) and \
            not name.startswith("__") and \
            not name.startswith("@"):
            globals_backup[name] = obj

    # 终于进入正题:重新载入修改过后的测试代码,拿到new module
    new_module = importlib.reload(old_module)
    # 拿到当前运行失败的test case函数名
    name = item.function.__name__
    # 注意:pytest的一个test case既可以是module里的一个function,也可以是test class里面的一个method,所以要分别处理
    if hasattr(new_module, name):
        # 如果是function就比较好办,直接从new module里拿出新的test case function,然后覆盖即可
        f = getattr(new_module, name)
        item.obj = f
    else:
        # 但如果是class method,我们就不得不遍历new module里的所有class,取出这个名字的method
        f = load_method_from_classes(new_module, name)
        # 然后需要做一个比较tricky的操作,覆盖的是`__code__`属性,而不能直接覆盖`obj`
        item.obj.__func__.__code__ = f.__code__

    # 最后我们还需要把old module里的全局变量复制到new module里面,因为test case里的逻辑有可能还会依赖这些数据
    for name, obj in globals_backup.items():
        if _is_empty(getattr(new_module, name)):
            setattr(new_module, name, obj)
    print("\nReloaded test module from [{}]".format(new_module.__file__))

经过各种corner case的修修补补之后,最终的完整代码实现在这里

work with ipdb

ipdbIPython 版的 pdb,提供语法高亮、自动补全等增强功能,使调试更高效,在需要的地方加上 import ipdb; ipdb.set_trace() 就可以设置断点。这个调试器配合我们的pytest插件就可以完美解决写test case时搞不清楚动态类型的难题,因为我们可以随时拿到一个ipython shell来检查当前scope里的各种对象到底该怎么调用,然后修改测试逻辑直到它能跑通为止。

comments powered by Disqus
Published:
2025-01-27
分类:
Tag: