如何从pytest-xdist节点获取数据
前言
应业务需求，需要用到pytest-xdist库作并行任务，为解决钩子函数重复调用问题，所以记录一下。
主要从博主发表的文章获得启发.
原文地址：https://korytkin.medium.com/how-to-get-data-from-pytest-xdist-nodes-2fbf2f0fe957（需要梯子）
GitHub: https://gist.github.com/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 -lvv tests -n 3
需要做一些hack来迫使它工作。但是首先，需要理解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
总结：只需要再我们的钩子函数中判断是否有workeroutput即可确定是否master
