Unit test conditional statements using environment variable

Question

I have a source file util.py and there I have below one, where day is an AWS lambda environment variable.

import os

day = os.getenv("DAY")

def get_day_code():
    if day == "Monday":
        return "M"
    if day == "Tuesday":
        return "T"
    if day == "Wednesday":
        return "W"

I need to write a unit test for this function. I tried using pytest paramaterise, but nothing works. Any help would be really appreciated.

Answer 1

Since os.getenv returns None, if no environment variable with the provided name is found, you don't need to mock that function. Instead, you can simply patch the global variable day during your tests.

unittest.mock.patch can be used in the form of a context manager. That way you can easily probe every branch of your function in a loop.

Assuming your code resides in a module named util, here is a working example test module:

from unittest import TestCase
from unittest.mock import patch

from . import util


class MyTestCase(TestCase):
    def test_get_day_code(self) -> None:
        day_map = {
            "Monday": "M",
            "Tuesday": "T",
            "Wednesday": "W",
        }
        for day, expected_output in day_map.items():
            with patch.object(util, "day", new=day):
                self.assertEqual(expected_output, util.get_day_code())

---

If you want to use pytest, you can also inject temporary environment variables using monkeypatch.setenv, but I would argue that there again directly patching the day variable is the easier route to take. The test should look essentially the same as demonstrated above.