前言

应业务需求，需要用到pytest-xdist库作并行任务，为解决钩子函数重复调用问题，所以记录一下。主要从博主发表的文章获得启发. 原文地址：https://korytkin.medium wangt.cc /how-to-get-data-from-pytest-xdist-nodes-2fbf2f0fe957（需要梯子） GitHub: https://gist.github wangt.cc /DKorytkin/8a186693af9a015abe89f6b874ca0795

简单pytest的插件功能实现

我们实现了一个简单的pytest插件，它可以通过运行并返回到占用大量RAM的前5个测试来测量每个测试的内存使用统计数据。

import collections
import os

import psutil
import pytest


LIMIT = 5


def pytest_configure(config):
    """
    Defined appropriate plugins selection in pytest_configure hook
    Parameters
    ----------
    config : _pytest.config.Config
    """
    plugin = MemoryUsage(config)
    config.pluginmanager.register(plugin)


def get_usage_memory():
    """
    Measures memory usage per Python process
    Returns
    -------
    memory_usage : float
    """
    process = psutil.Process(os.getpid())
    memory_use = process.memory_info()
    return memory_use.rss / 1024  # to Kb


class MemoryUsage:
    def __init__(self, config):
        """
        Defined appropriate plugins selection in pytest_configure hook
        Parameters
        ----------
        config : _pytest.config.Config
        """
        self.config = config
        self.stats = collections.defaultdict(dict)

    def pytest_runtest_setup(self, item):
        """Record maxrss for pre-setup."""
        self.stats[item.nodeid]["setup"] = get_usage_memory()

    @pytest.hookimpl(hookwrapper=True)
    def pytest_runtest_call(self, item):
        """
        Track test memory
        Parameters
        ----------
        item : _pytest.main.Item
        """
        start = get_usage_memory()
        yield
        end = get_usage_memory()
        self.stats[item.nodeid]["diff"] = end - start
        self.stats[item.nodeid]["end"] = end
        self.stats[item.nodeid]["start"] = start

    def pytest_terminal_summary(self, terminalreporter):
        tr = terminalreporter
        if self.stats:
            tr._tw.sep("=", "TOP {} tests which took most RAM".format(LIMIT), yellow=True)
            stats = sorted(self.stats.items(), key=lambda x: x[-1]["diff"], reverse=True)
            for test_name, info in stats[:LIMIT]:
                line = "setup({}Kb) usage ({}Kb) - {}".format(info["setup"], info["diff"], test_name)
                tr._tw.line(line)

看起来我们的插件工作正常🔥

如何从pytest-xdist节点获取数据

但是过了一段时间，当我们的项目变得庞大时，需要考虑并行化，首先想到的是pytest-xdist，它将是解决我们问题的好工具。但是出乎意料的是我们的插件不工作😵

pytest -lvv tests -n 3

需要做一些hack来迫使它工作。但是首先，需要理解pytest-xdist是如何工作的

如何从pytest-xdist节点获取数据

pytest-xdist工作原理

pytest-xdist需要并行运行测试，当你运行pytest -n 4 tests/backend/unit where -n 4 number of nodes will run for testing 这意味着，它将运行5 个独立的 python 进程：

master
gw0
gw1
gw2
gw3

主节点不运行任何测试，只是通过一小部分消息与节点通信，例如:

workerready （当节点成功启动时）
collectionstart （集合启动）
collectionfinish (收集完成)
runtest_protocol_complete (已完成的单项测试)
等等……

我试图展示一个简单的图表，说明它一般是如何工作的

workerinput (数据从主节点发送到节点)
workeroutput (数据从节点发送到主节点)

然后，当我们已经知道它是如何工作的，我们可以做一些事情，并修复我们的插件😉 这里的主要思想是使用workeroutput。节点将workeroutput发送到主节点，我们可以在pytest_sessionfinish钩子中将我们的信息添加到这个字典中，这个钩子也调用每个节点和主节点(最后一个)，为了了解我们在哪里，我们可以检查配置中的workeroutput属性，如果它不存在，我们在主节点中。

SHARED_MEMORY_USAGE_INFO = "memory_usage"


def is_master(config):    
    """
    True if the code running the given pytest.config object is     
    running in a xdist master    node or not running xdist at all.        
    """
    return not hasattr(config, 'workerinput')


@pytest.hookimpl(hookwrapper=True, trylast=True)
def pytest_sessionfinish(self, session, exitstatus):
    """
    Dump memory usage statistics to `workeroutput`
    Executed once per node if with xdist and will get from mater node
    Parameters
    ----------
    session : _pytest.Session
    exitstatus : int
    """
    yield
    if not self.is_master:
        self.config.workeroutput[SHARED_MEMORY_USAGE_INFO] = self.stats

之后，我们可以将所有从节点收到的数据合并到主节点的单个Dict中，这种能力是在pytest_testnodedown钩子中允许的，当节点中的所有测试完成时，该钩子在主节点中被调用一次。

def pytest_testnodedown(self, node, error):
    """
    Get statistic about memory usage for test cases from xdist nodes
    and merge to master stats
    """
    node_stats = node.workeroutput[SHARED_MEMORY_USAGE_INFO]
    self.stats.update(node_stats)

这对我们来说已经足够了，我们的插件再次工作正常，并且已经支持pytest-xdist

如何从pytest-xdist节点获取数据