Temp environment variables for pytest

The Problem

Ever had a Python project that uses a tool or package that is configured by environment variables to authenticates with some service? When that is the case, we often have our project locally configured to work against the actual “working” service environment, usually using the credentials of our personal user on that service — for example, using the S3 of our AWS research account.

However, when we write integration tests for the same project, we might want them to work against an actual environment of that service (and not against a mock of it), but still not against a production environment. If we use a CI service (like Travis or Circle CI), we can set up credentials for a different user — or even for a different environment — just for test runs, but we might also want to be able to run test locally against the testing user/environment without interfering with “normal”/actual runs of the code, meant to run against production/working environments.

If that is the case, and you’re using pytest to test your project, you might find the following solution helpful.

I found it useful when wanting to run local (laptop-bound) tests for a project that uses a Databricks-hosted mlflow server for experiment tracking — authenticating using environment variables used by the databricks-cli Python package — against a dedicated test server, without messing up my local runs of actual experiments of the same project writing to the research mlflow server.

The Solution

The solution is made up of two Python files, both to be added to your pytest tests directory (usually tests in the root of your project):

1. conftest.py — A standard way to configure advanced pytest features such as hooks, fixtures (our use case) and plugins.
2. temp_env_var.py— A Python file to list the temporary, test-only, environment variables to set up.

You tests directory should thus end up looking something like this:

tests\
  __init__.py
  conftest.py
  temp_env_var.py
  test_basic_tests.py
  test_advanced_functionality.py
  ...

First, set up the environment variables you want to add and/or remove for the duration of the test in temp_env_var.py , using the appropriate global variables defined there.

For example, I used this to set up credentials for my test Databricks account, also making sure an alternative authentication method — via a token — is nullified for the duration of the tests, without messing up the credentials for the production account, stored in the same environment variables:

Now, here’s how conftest.py uses them:

Since this is a quick-and-dirty gist/tutorial, I won’t explain what pytest fixtures are, but basically:

• The combination of scope="session" and autouse=True means this code runs only once for the whole test suit.
• Everything before the yield will be performed before all tests are run.
• Everything after the yield will be performed after all tests are run.

This is a common setup/tear-down pattern for pytest.

As you can see, the code simply imports the environment variables to setup from your untracked variables file, gracefully creating an empty dict and an empty list if the import fails (for example, on our remote CI server which also runs these tests, but has these variables set up normally for the project).

It then keeps a copy of os.environ, the mapping Python uses to represent the string environment (and which your third-party dependency uses to read environment variables), and updates this mapping with the temporary variables, also removing forbidden ones.

Finally, when all tests finish running, it restores the pre-test environment.

Important: Add */temp_env_var.py to your .gitignore file, to ensure test credentials are never committed into your git repository!

That’s it. I hope some of you will find this useful. Let me know if I’ve made a mistake, or if you’ve got a better way to achieve this. :)

If different parts of your test suit require different values for the same environment variables, I suggest you take the time to understand pytest fixtures a bit more in-depth, and then use the ability to define them for smaller scopes (namely, function, class, module and package) to expand the above approach to handle such cases.

Sources

• I used this StackOverflow thread for the code that manages the temporary environment variable setup.