如何全局种子 np.random.default_rng 进行单元测试

Question

numpy 创建随机数的推荐方法是创建一个np.random.Generator像这样的

import numpy as np

def foo():
    # Some more complex logic here, this is the top level method that creates the rng
    rng = np.random.default_rng()
    return rng.random()

现在假设我正在为我的代码库编写测试，并且我需要为 rng 提供种子以获得可重现的结果。

是否可以告诉 numpy 每次都使用相同的种子，无论在哪里default_rng()调用？这基本上是 的旧行为np.random.seed()。我需要这个的原因是因为我有很多这样的测试，并且必须模拟调用default_rng以对每个测试使用种子，因为在 pytest 中，您必须在使用某些东西的位置而不是定义它的位置进行模拟。因此，像这个答案一样在全球范围内嘲笑它是行不通的。

使用旧方法，可以定义一个固定装置，在 conftest.py 中自动设置每个测试的种子，如下所示：

# conftest.py

import pytest
import numpy as np

@pytest.fixture(autouse=True)
def set_random_seed():
    # seeds any random state in the tests, regardless where is is defined
    np.random.seed(0)

# test_foo.py

def test_foo():
    assert np.isclose(foo(), 0.84123412)  # That's not the right number, just an example

随着新的使用方式default_rng，这似乎不再可能了。相反，我需要在每个需要播种 rng 的测试模块中放置这样的夹具。

# inside test_foo.py, but also every other test file

import pytest
from unittest import mock
import numpy as np


@pytest.fixture()
def seed_default_rng():
    seeded_rng = np.random.default_rng(seed=0)
    with mock.patch("module.containing.foo.np.random.default_rng") as mocked:
        mocked.return_value = seeded_rng
        yield 

def test_foo(seed_default_rng):
    assert np.isclose(foo(), 0.84123412)

Answer 1

我想出的最好的办法是在 conftest.py 中有一个可参数化的固定装置，如下所示

# conftest.py
import pytest
from unittest import mock
import numpy as np


@pytest.fixture
def seed_default_rng(request):
    seeded_rng = np.random.default_rng(seed=0)
    mock_location = request.node.get_closest_marker("rng_location").args[0]
    with mock.patch(f"{mock_location}.np.random.default_rng") as mocked:
        mocked.return_value = seeded_rng
        yield

然后可以在每个测试中使用它，如下所示：

# test_foo.py
import pytest
from module.containing.foo import foo

@pytest.mark.rng_location("module.containing.foo")
def test_foo(seed_default_rng):
    assert np.isclose(foo(), 0.84123412)  # just an example number

虽然还是没有以前那么方便，但是只需要在每个测试中添加标记即可，而不用模拟 default_rng 方法。