Browse Source

tests: helper function to temporarily change the environment

keep-around/8cd9a86ebf4bfe8f134748e8ac721cfaba15bbd4
singuliere 3 years ago
parent
commit
fb2f835f3a
No known key found for this signature in database GPG Key ID: 900857755EF189C2
  1. 35
      tests/modified_environ.py
  2. 46
      tests/test_modified_environ.py

35
tests/modified_environ.py

@ -0,0 +1,35 @@
# coding: utf-8
# Copy/pasted from https://github.com/laurent-laporte-pro/stackoverflow-q2059482/blob/master/demo/environ_ctx.py # noqa
import contextlib
import os
@contextlib.contextmanager
def modified_environ(*remove, **update):
"""
Temporarily updates the ``os.environ`` dictionary in-place.
The ``os.environ`` dictionary is updated in-place so that the modification
is sure to work in all situations.
:param remove: Environment variables to remove.
:param update: Dictionary of environment variables and values to add/update.
"""
env = os.environ
update = update or {}
remove = remove or []
# List of environment variables being updated or removed.
stomped = (set(update.keys()) | set(remove)) & set(env.keys())
# Environment variables and values to restore on exit.
update_after = {k: env[k] for k in stomped}
# Environment variables and values to remove on exit.
remove_after = frozenset(k for k in update if k not in env)
try:
env.update(update)
[env.pop(k, None) for k in remove]
yield
finally:
env.update(update_after)
[env.pop(k) for k in remove_after]

46
tests/test_modified_environ.py

@ -0,0 +1,46 @@
# coding: utf-8
# copy/pasted from https://github.com/laurent-laporte-pro/stackoverflow-q2059482/blob/master/tests/test_environ_ctx.py # noqa
import os
from tests.modified_environ import modified_environ
def setup_method(test_method):
os.environ.pop('MODIFIED_ENVIRON', None)
def teardown_method(test_method):
os.environ.pop('MODIFIED_ENVIRON', None)
def test_modified_environ__no_args():
with modified_environ():
pass
def test_modified_environ__inserted():
with modified_environ(MODIFIED_ENVIRON="inserted"):
assert os.environ['MODIFIED_ENVIRON'] == "inserted"
assert 'MODIFIED_ENVIRON' not in os.environ
def test_modified_environ__updated():
os.environ['MODIFIED_ENVIRON'] = "value"
with modified_environ(MODIFIED_ENVIRON="updated"):
assert os.environ['MODIFIED_ENVIRON'] == "updated"
assert os.environ['MODIFIED_ENVIRON'] == "value"
def test_modified_environ__deleted():
os.environ['MODIFIED_ENVIRON'] = "value"
with modified_environ('MODIFIED_ENVIRON'):
assert 'MODIFIED_ENVIRON' not in os.environ
assert os.environ['MODIFIED_ENVIRON'] == "value"
def test_modified_environ__deleted_missing():
with modified_environ('MODIFIED_ENVIRON'):
assert 'MODIFIED_ENVIRON' not in os.environ
assert os.environ['MODIFIED_ENVIRON'] == "value"
Loading…
Cancel
Save