使用 pytest fixture 下载一次数据,但每次测试都使用新副本

Use a pytest fixture to download data once, but use new copies for each test

我有下载一组图像的测试代码,对它们进行一些处理,并断言处理按预期进行:

@pytest.fixture
def image_dir(tmp_path):
    test_imgs = [
        # ... list of img URLs
    ]
    for idx, im_url in enumerate(test_imgs):
        urllib.request.urlretrieve(im_url, tmp_path / f"{idx}.png")
    yield tmp_path

def test_op_A(image_dir: Path):
    for im_path in image_dir.iterdir():
        # load the image
        # modify the image
        # save the image back to disk

    # assert that the modification worked as expected

def test_op_B(image_dir: Path):
    for im_path in image_dir.iterdir():
        # load the image
        # modify the image
        # save the image back to disk

    # assert that the modification worked as expected

# ... more tests with a similar format

这可行,但速度非常慢。我怀疑这是因为每次测试都重新下载图像。

有没有一种干净的方法来创建一次临时目录,缓存它,并为每个测试使用目录的副本?这样每个测试都可以根据需要修改图像,而不会影响其他测试,同时保持性能。

因此,实现此目标的可能解决方案是使用 pytest_sessionstart and pytest_sessionfinish。我们还将使用夹具来复制文件。

我们希望实现的流程的一般分解如下:

  1. 在测试收集之前下载必要的文件
  2. 让 fixture 将必要的文件复制到临时位置
  3. 测试完成后,从缓存位置删除文件

将以下两个挂钩放在您的 conftest.py 测试所在的根目录中。

from pathlib import Path
from tempfile import TemporaryDirectory

def pytest_sessionstart(session):
    test_imgs = [
        # ... list of img URLs
    ]
    td = TemporaryDirectory()
    tmp_path = Path(td.name)

    for idx, im_url in enumerate(test_imgs):
        urllib.request.urlretrieve(im_url, tmp_path / f"{idx}.png")

    session.__IMAGE_CACHE = tmp_path


def pytest_sessionfinish(session, exitstatus):
    # remove the cached images
    session.__IMAGE_CACHE.cleanup()

现在我们已经将图像缓存在一个位置,我们可以让我们的 fixture 为每个测试复制它们,而不必下载它们。

from shutil import copytree

@pytest.fixture
def image_dir(tmp_path, request):
    session = request.session
    # copy the data from our cache to the temp location for the test
    copytree(session.__IMAGE_CACHE, tmp_path)
    yield tmp_path