如何在pytest中设置环境变量

Question

我有一个使用环境变量的lambda handler。我如何使用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):
...

- Antarr Byrd

---

在调用pytest时设置变量怎么样？'BUCKET=foobar pytest...' - jkr

Answer 1

你在猴子补丁运行之前就遇到了失败。当导入runner模块时，环境变量的加载会发生。

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

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

如果您无法修改该文件，则事情会更加复杂。

我尝试创建一个全局固件，使用monkeypatch对环境进行修改并在所有测试加载之前仅加载一次模块，但收到了pytest错误，指出在session级别的固件中使用了function level fixtures。这很有道理，因为monkeypatch并不打算长期伪造东西。您可以将模块加载放入测试之后进行monkeypatch，但这会生成大量样板代码。

最终成功的方法是创建一个fixture来代替导入类。该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 加载时创建的环境变量。如果你希望每个测试都使用同一个值，那么你只需要在加载任何测试模块之前将所需的值添加到其中即可。在 conftest.py 的顶部加上 os.environ['BUCKET'] = 'test-bucket' 就可以为整个测试会话设置环境变量。只要模块的第一次导入发生在此之后，就不会出现关键错误。这种方法的一个明显缺点是，除非你知道要查看 conftest.py 或者在代码中进行 grep，否则很难确定在故障排除时环境变量从哪里被设置。

- sphennings

Answer 2

当我在寻找pytest中模拟环境变量的方法时，我看到了你的问题。

我发现Adam Johnson提供了一种使用pytest fixtures和python unittest包中的mock来模拟os.environ的解决方案。它可以归结为以下内容：

import os
from unittest import mock

import pytest

@pytest.fixture(scope='class')
def mock_settings_env_vars(scope='class'):
    with mock.patch.dict(os.environ, {"FROBNICATION_COLOUR": "ROUGE"}):
        yield

class TestEnvMocking:
    def test_frobnication_colour_not_yet_available(self):
        assert "FROBNICATION_COLOUR" not in os.environ
    def test_frobnication_colour(self, mock_settings_env_vars):
        assert os.environ["FROBNICATION_COLOUR"] == "ROUGE"
    def test_frobnication_colour_sideeffect_still_there(self):
        assert os.environ["FROBNICATION_COLOUR"] == "ROUGE"

class TestEnvMockingGone:
    def test_frobnication_colour_not_available(self):
        assert "FROBNICATION_COLOUR" not in os.environ

模拟夹具的作用域为类级别，因此模拟对象将从第一次使用开始一直保留到测试类完成。

- itsafire