如何在pytest中设置环境变量

Question

我有一个使用环境变量的巴巴处理程序。我如何使用 pytest 设置该值？我收到错误

tests/test_kinesis.py:3: in <module>
    from runner import kinesis
runner/kinesis.py:6: in <module>
    DATA_ENGINEERING_BUCKET = os.environ["BUCKET"]
../../../../../.pyenv/versions/3.8.8/lib/python3.8/os.py:675: in __getitem__
    raise KeyError(key) from None
E   KeyError: 'BUCKET'
7:03

我尝试在测试中这样设置

tests/test_kinesis.py:3: in <module>
    from runner import kinesis
runner/kinesis.py:6: in <module>
    DATA_ENGINEERING_BUCKET = os.environ["BUCKET"]
../../../../../.pyenv/versions/3.8.8/lib/python3.8/os.py:675: in __getitem__
    raise KeyError(key) from None
E   KeyError: 'BUCKET'
7:03

class TestHandler(unittest.TestCase):
    @mock_s3
    @mock_lambda
    def test_handler(monkeypatch):
        monkeypatch.setenv("BUCKET", "test-bucket")
        actual = kinesis.handler(kinesis_stream_event, "")
        expected = {"statusCode": 200, "body": "OK"}
        assert actual == expected

Answer 1

在你的猴子补丁能够运行之前你就遇到了失败。首次导入运行程序模块时将加载环境变量。

如果这是您拥有的模块，如果未设置 DATA_ENGINEERING_BUCKET，我建议修改代码以使用默认值。然后您可以通过调用在运行时将其值修改为您想要的任何值module.DATA_ENGINEERING_BUCKET = "my_bucket"。

DATA_ENGINEERING_BUCKET = os.environ.get("BUCKET", default="default_bucket")

如果您无法修改该文件，那么事情就会变得更加复杂。

我考虑创建一个全局固定装置，在加载任何测试之前对环境进行猴子修补并加载模块一次，并收到有关在会话级固定装置中使用功能级固定装置的 pytest 错误。这是有道理的，monkeypatch 实际上并不是为了长期伪造东西。您可以在猴子补丁之后将模块加载粘贴到测试中，但这会生成大量样板文件。

最终创建了一个固定装置来提供该类而不是导入它。固定装置；设置os.environ为所需的值，加载模块，将 os.environ 重置为其原始值，然后生成模块。任何需要此模块的测试都可以请求夹具在其范围内访问它。需要注意的是，因为测试文件是在运行固定装置之前导入的，所以任何不使用固定装置并导入模块的测试文件通常都会引发 KeyError 并导致 pytest 在运行任何测试之前崩溃。

测试.py

import os, pytest

@pytest.fixture(scope='session')
def kinesis():
    old_environ = os.environ
    os.environ = {'BUCKET': 'test-bucket'}
    import kinesis
    os.environ = old_environ
    yield kinesis

测试.py

# Do NOT import kinesis in any test file. Rely on the fixture.
class TestHandler(unittest.TestCase):
    @mock_s3
    @mock_lambda
    def test_handler(kinesis):
        actual = kinesis.handler(kinesis_stream_event, "")
        expected = {"statusCode": 200, "body": "OK"}
        assert actual == expected

一种可能更简单的方法

os.environ是操作系统首次加载时创建的环境变量字典。如果您希望每个测试都有一个值，那么您只需在加载任何测试模块之前添加所需的值即可。如果您将其放在os.environ['BUCKET'] = 'test-bucket'conftest.py 的顶部，您将为测试会话的其余部分设置环境变量。然后，只要之后发生模块的第一次导入，就不会出现关键错误。这种方法的一大缺点是，除非您知道查看 conftest.py 或 grep 代码，否则在故障排除时很难确定环境变量的设置位置。