如何在pytest中设置环境变量

Question

我有一个使用环境变量的lamba处理程序。如何使用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

我试着在测试中这样设置

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

DATA_ENGINEERING_BUCKET = os.environ["BUCKET"]


def handler(event, context):
...

Answer 1

在你的monkeypatch能够运行之前，你就已经失败了。环境变量的加载将在第一次导入runner模块时发生。

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

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

如果你不能修改这个文件，那么事情就更复杂了。

我考虑创建一个全局fixture，在加载任何测试之前，monkeypatches环境和加载模块一次，并收到关于在会话级fixture中使用函数级fixture的pytest错误。这就说得通了，monkeypatch并不是用来长期伪造东西的。你可以在monkeypatch之后把模块加载到你的测试中，但是这会产生很多样板。

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

conftest.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

tests.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中查找代码，否则很难在进行故障排除时确定环境变量设置的位置。