如何在pytest中设置环境变量

Question

I have a lamba handler that uses an environment variable. How can I set that value using pytest. I'm getting the error

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

I tried setting in the test like this

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

You're getting the failure before your monkeypatch is able to run. The loading of the environment variable will happen when the runner module is first imported.

If this is a module you own, I'd recommend modifying the code to use a default value if DATA_ENGINEERING_BUCKET isn't set. Then you can modify it's value to whatever you want at runtime by calling module.DATA_ENGINEERING_BUCKET = "my_bucket" .

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

If you can't modify that file then things are more complicated.

I looked into creating a global fixture that monkeypatches the environment and loads the module once, before any tests load and received a pytest error about using function level fixtures within a session level fixture. Which makes sense monkeypatch really isn't intended to fake things long term. You can stick the module load into your test after the monkeypatch but that will generate a lot of boilerplate.

What eventually worked creating a fixture that will provide the class in lieu of importing it. The fixture; sets os.environ to the desired value, loads the module, resets os.environ to it's origional value then yields the module. Any tests that need this module can request the fixture to have access to it within their scope. A word of caution, because test files are imported before fixtures are run any test files that don't use the fixture and import the module normally will raise a KeyError and cause pytest to crash before running any tests.

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

A potentially simpler method

os.environ is a dictionary of environment variables that is created when os first loads. If you want a single value for every test then you just need to add the value you want to it before loading any test modules. If you put os.environ['BUCKET'] = 'test-bucket' at the top of conftest.py you will set the environment variable for the rest of the test session. Then as long as the first import of the module happens afterwards you won't have a key error. The big downside to this approach is that unless you know to look in conftest.py or grep the code it will be difficult to determine where the environment variable is getting set when troubleshooting.