How to unit-test a decorator which uses environment variable?

Question

I created a decorator to allow functions be be run in specific environments only:

def accepted_environments(*envs):
    """
    The decorated function can be executed only in specified envs
    """
    def my_decorator(func_to_be_decorated):
        def wrapper():
            if settings.ENV_NAME not in envs:
                raise EnvironmentException
            return func_to_be_decorated()
        return wrapper
    return my_decorator

# Usage example
@accepted_environments('local', 'prod')
def hello():
    print("hello")

That seems to work, but I'd like to unit test it. The problem is: my tests are potentialy run in every environment (local, staging, prod). Plus, I think it's not safe my tests are able to change environment variables.

Should I "mock" this behaviour? How would you do? Thanks!

Answer 1

Use mock to override the value of settings.ENV_NAME for a test.

def test_not_in_dev(self):
    with mock.patch.dict(settings.__dict__, ENV_NAME="dev"):
        self.assertRaises(EnvironmentException, hello)