如何让Pytest测试在调用未await的async函数时自动失败？

原方案的问题 #

你提供的ChatGPT代码无法生效，因为它仅检查了测试函数自身返回的协程，完全没有处理测试函数内部调用async函数生成的协程对象。对于test_await内部的func()调用，该装饰器根本无法感知到，自然不会触发失败。

可行方案 #

下面提供两种实现思路，分别基于静态代码分析和运行时跟踪，都能满足你的需求。

方案一：静态代码分析（AST） #

通过解析测试函数的抽象语法树（AST），提前检测所有async函数调用是否被await包裹。这种方式在测试执行前就能发现问题，效率较高。

import ast
import inspect
import functools

def checkawait(test_func):
    @functools.wraps(test_func)
    async def wrapper(*args, **kwargs):
        # 获取测试函数的源代码
        try:
            source = inspect.getsource(test_func)
        except OSError:
            # 处理无法获取源代码的情况（如动态生成的函数）
            return await test_func(*args, **kwargs)
        
        # 解析AST
        tree = ast.parse(source)
        
        # 遍历AST节点，检查未被await的async函数调用
        class CoroutineCallChecker(ast.NodeVisitor):
            def visit_Call(self, node):
                # 判断当前调用是否被await包裹
                is_awaited = isinstance(getattr(node, 'parent', None), ast.Await)
                
                # 尝试获取被调用的函数对象
                func_obj = None
                try:
                    if isinstance(node.func, ast.Name):
                        # 处理直接调用全局/局部函数的情况（如func()）
                        func_name = node.func.id
                        func_obj = test_func.__globals__.get(func_name)
                        if not func_obj:
                            func_obj = test_func.__closure__[0].cell_contents.get(func_name) if test_func.__closure__ else None
                    elif isinstance(node.func, ast.Attribute):
                        # 处理属性调用的情况（如obj.method()）
                        if isinstance(node.func.value, ast.Name):
                            obj_name = node.func.value.id
                            obj = test_func.__globals__.get(obj_name)
                            if obj:
                                func_obj = getattr(obj, node.func.attr, None)
                except (AttributeError, IndexError, KeyError):
                    pass
                
                # 检测到未被await的async函数调用
                if func_obj and inspect.iscoroutinefunction(func_obj) and not is_awaited:
                    raise AssertionError(f"Async function '{func_obj.__name__}' was called without await")
                
                self.generic_visit(node)
        
        # 给所有AST节点添加parent属性，用于判断是否被await包裹
        for node in ast.walk(tree):
            for child in ast.iter_child_nodes(node):
                child.parent = node
        
        # 执行检查
        checker = CoroutineCallChecker()
        checker.visit(tree)
        
        # 执行原测试函数
        return await test_func(*args, **kwargs)
    
    return wrapper

优缺点：

• 优点：无需运行测试逻辑即可提前发现问题，性能开销小；错误信息清晰，直接指出未被await的函数名。
• 缺点：无法处理动态生成的函数、闭包内的函数或间接引用的函数（如通过变量调用async函数）；对复杂语法场景的支持有限。

方案二：运行时跟踪（sys.settrace） #

通过Python的sys.settrace钩子在测试执行期间跟踪所有函数调用，记录生成的协程对象，并检查这些协程是否被await处理。这种方式能覆盖更多动态场景。

import sys
import asyncio
import inspect
import functools

def checkawait(test_func):
    @functools.wraps(test_func)
    async def wrapper(*args, **kwargs):
        created_coroutines = set()
        awaited_coroutines = set()
        
        def trace_handler(frame, event, arg):
            if event == 'call':
                # 检查当前调用的是否是async函数
                func_name = frame.f_code.co_name
                func_obj = frame.f_globals.get(func_name)
                if func_obj and inspect.iscoroutinefunction(func_obj):
                    # async函数调用会返回协程对象，arg即为该对象
                    if inspect.iscoroutine(arg):
                        created_coroutines.add(arg)
            elif event == 'yield':
                # 当await协程时，arg通常是被await的协程对象
                if inspect.iscoroutine(arg):
                    awaited_coroutines.add(arg)
            return trace_handler
        
        # 设置跟踪钩子
        original_trace = sys.settrace(trace_handler)
        try:
            # 执行原测试函数
            result = await test_func(*args, **kwargs)
        finally:
            # 恢复原跟踪钩子
            sys.settrace(original_trace)
        
        # 检查未被await的协程
        unawaited = created_coroutines - awaited_coroutines
        if unawaited:
            raise AssertionError(f"Found {len(unawaited)} unawaited coroutine(s)")
        
        return result
    
    return wrapper

优缺点：

• 优点：能检测动态生成的函数、间接调用的async函数；覆盖更多运行时场景。
• 缺点：运行时跟踪会带来一定性能开销；无法区分“故意不await（如创建后台任务）”和“忘记await”的情况，需要额外逻辑处理例外。

验证示例 #

使用任意一种装饰器后，你的测试代码会触发失败：

async def func():
    return 42

@checkawait
async def test_await():
    func()  # 未加await，测试会抛出AssertionError

执行该测试时，会直接报错并终止测试，符合你的需求。

内容的提问来源于stack exchange，提问作者ychiucco