Writing Unit Tests for Python Code (unittest & doctest)

In Python, there are several tools to help you write, organize, run, and automate your unit tests. In the Python standard library, you’ll find two of these tools:

doctest
unittest (inspired by Junit)

Python’s doctest module is a lightweight testing framework that provides quick and straightforward test automation. It can read the test cases from your project’s documentation and your code’s docstrings.

The unittest package is also a testing framework. However, it provides a more complete solution than doctest.

Test case: An individual unit of testing. It examines the output for a given input set.
Test suite: A collection of test cases, test suites, or both. They’re grouped and executed as a whole.
Test fixture: A group of actions required to set up an environment for testing. It also includes the teardown processes after the tests run.
Test runner: A component that handles the execution of tests and communicates the results to the user.

# age.py

def categorize_by_age(age):
    if 0 <= age <= 9:
        return "Child"
    elif 9 < age <= 18:
        return "Adolescent"
    elif 18 < age <= 65:
        return "Adult"
    elif 65 < age <= 150:
        return "Golden age"
    else:
        return f"Invalid age: {age}"

# test_age.py

import unittest

from age import categorize_by_age

class TestCategorizeByAge(unittest.TestCase):
    def test_child(self):
        self.assertEqual(categorize_by_age(5), "Child")

    def test_adolescent(self):
        self.assertEqual(categorize_by_age(15), "Adolescent")

    def test_adult(self):
        self.assertEqual(categorize_by_age(30), "Adult")

    def test_golden_age(self):
        self.assertEqual(categorize_by_age(70), "Golden age")

    def test_negative_age(self):
        self.assertEqual(categorize_by_age(-1), "Invalid age: -1")

    def test_too_old(self):
        self.assertEqual(categorize_by_age(151), "Invalid age: 151")

Running unittest Tests

Make the test module executable
Use the command-line interface of unittest

1. Make the test module executable

# test_age.py

#...

if __name__ == "__main__":
    unittest.main()


# now you can directly run this file
# python test_age.py

Among other arguments, the main() function takes the verbosity one. With this argument, you can tweak the output’s verbosity, which has three possible values:

0 for quiet
1 for normal
2 for detailed
Go ahead and update the call to main() as in the following:

# ...

if __name__ == "__main__":
    unittest.main(verbosity=2)

# output

.
.
.
python test_age.py
test_adolescent (__main__.TestCategorizeByAge.test_adolescent) ... ok
test_adult (__main__.TestCategorizeByAge.test_adult) ... ok

If you want to make the detailed output more descriptive, then you can add docstrings to your tests like in the following code snippet:


# ...

class TestCategorizeByAge(unittest.TestCase):
    def test_child(self):
        """Test for 'Child'"""
        self.assertEqual(categorize_by_age(5), "Child")

    def test_adolescent(self):
        """Test for 'Adolescent'"""
        self.assertEqual(categorize_by_age(15), "Adolescent")

    def test_adult(self):
        """Test for 'Adult'"""
        self.assertEqual(categorize_by_age(30), "Adult")

# ...

# output
python test_age.py
test_adolescent (__main__.TestCategorizeByAge.test_adolescent)
Test for 'Adolescent' ... ok
test_adult (__main__.TestCategorizeByAge.test_adult)
Test for 'Adult' ... ok
test_boundary_adolescent_adult (__main__.TestCategorizeByAge.test_boundary_adolescent_adult)
Test for boundary between 'Adolescent' and 'Adult' ... ok

Skipping Tests

import sys
import unittest

class SkipTestExample(unittest.TestCase):
    @unittest.skip("Unconditionally skipped test")
    def test_unimportant(self):
        self.fail("The test should be skipped")

    @unittest.skipIf(sys.version_info < (3, 12), "Requires Python >= 3.12")
    def test_using_calendar_constants(self):
        import calendar

        self.assertEqual(calendar.Month(10), calendar.OCTOBER)

    @unittest.skipUnless(sys.platform.startswith("win"), "Requires Windows")
    def test_windows_support(self):
        from ctypes import WinDLL, windll

        self.assertIsInstance(windll.kernel32, WinDLL)

if __name__ == "__main__":
    unittest.main(verbosity=2)

Why do we need Sub tests?

class TestIsEven(unittest.TestCase):
    def test_even_number(self):
        for number in [2, 4, 6, -8, -10, -12]:
            self.assertEqual(is_even(number), True)

This will work fine, and the test will fail if any of the assertions fail. However, using subTest provides better granularity and diagnostics when running tests. Here’s why it’s useful:

✅ Benefits of `subTest`

Isolated Failures: Each iteration is treated as a separate subtest. If one number fails, the others still run, and you get a report for each.
Clearer Output: The test runner shows which specific input caused the failure, making debugging easier.
Improved Test Reporting: Especially useful when testing multiple inputs or edge cases.

Without `subTest`

If one assertion fails, the loop stops, and you don’t get feedback on the remaining values.

Example Output Comparison

With subTest:

FAIL: test_even_number (number=-10)
FAIL: test_even_number (number=-12)

Without subTest:

FAIL: test_even_number

Available Assert Methods

Comparing Values
Comparing the result of a code unit with the expected value is a common way to check whether the unit works okay. The TestCase class defines a rich set of methods that allows you to do this type of check:

Method Comparison
.assertEqual(a, b) a == b
.assertNotEqual(a, b) a != b
.assertTrue(x) bool(x) is True
.assertFalse(x) bool(x) is False

Comparing Objects by Their Identity
TestCase also implements methods that are related to the identity of objects. In Python’s CPython implementation, an object’s identity is the memory address where the object lives. This identity is a unique identifier that distinguishes one object from another.

An object’s identity is a read-only property, which means that you can’t change an object’s identity once you’ve created the object. To check an object’s identity, you’ll use the is and is not operators.

Here are a few assert methods that help you check for an object’s identity:

Method Comparison
.assertIs(a, b) a is b
.assertIsNot(a, b) a is not b
.assertIsNone(x) x is None
.assertIsNotNone(x) x is not None

Comparing Collections
Another common need when writing tests is to compare collections, such as lists, tuples, strings, dictionaries, and sets. The TestCase class also has shortcut methods for these types of comparisons. Here’s a summary of those methods:

Method Comparison
.assertSequenceEqual(a, b) Equality of two sequences
.assertMultiLineEqual(a, b) Equality of two strings
.assertListEqual(a, b) Equality of two lists
.assertTupleEqual(a, b) Equality of two tuples
.assertDictEqual(a, b) Equality of two dictionaries
.assertSetEqual(a, b) Equality of two sets

Using unittest From the Command Line

unittest test_age.py

you can also add modules to test

Grouping Your Tests With the TestSuite Class

import unittest

from calculations import (
    add,
    divide,
    mean,
    median,
    mode,
    multiply,
    subtract,
)

class TestArithmeticOperations(unittest.TestCase):
    def test_add(self):
        self.assertEqual(add(10, 5), 15)
        self.assertEqual(add(-1, 1), 0)

    def test_subtract(self):
        self.assertEqual(subtract(10, 5), 5)
        self.assertEqual(subtract(-1, 1), -2)

    def test_multiply(self):
        self.assertEqual(multiply(10, 5), 50)
        self.assertEqual(multiply(-1, 1), -1)

    def test_divide(self):
        self.assertEqual(divide(10, 5), 2)
        self.assertEqual(divide(-1, 1), -1)
        with self.assertRaises(ZeroDivisionError):
            divide(10, 0)

class TestStatisticalOperations(unittest.TestCase):
    def test_mean(self):
        self.assertEqual(mean([1, 2, 3, 4, 5, 6]), 3.5)

    def test_median_odd(self):
        self.assertEqual(median([1, 3, 3, 6, 7, 8, 9]), 6)

    def test_median_even(self):
        self.assertEqual(median([1, 2, 3, 4, 5, 6, 8, 9]), 4.5)

    def test_median_unsorted(self):
        self.assertEqual(median([7, 1, 3, 3, 2, 6]), 3)

    def test_mode_single(self):
        self.assertEqual(mode([1, 2, 2, 3, 4, 4, 4, 5]), [4])

    def test_mode_multiple(self):
        self.assertEqual(set(mode([1, 1, 2, 3, 4, 4, 5, 5])), {1, 4, 5})

if __name__ == "__main__":
    unittest.main()

These tests work as expected. Suppose you need a way to run the arithmetic and statistical tests separately. In this case, you can create test suites. In the following sections, you’ll learn how to do that.

# ...

# in the same way, we can make a suite for statistical

def make_suite():
    arithmetic_tests = [
        TestArithmeticOperations("test_add"),
        TestArithmeticOperations("test_subtract"),
        TestArithmeticOperations("test_multiply"),
        TestArithmeticOperations("test_divide"),
    ]
    return unittest.TestSuite(tests=arithmetic_tests)

if __name__ == "__main__":
    suite = make_suite()
    runner = unittest.TextTestRunner(verbosity=2)
    runner.run(suite)
# output

-> python test_calculations.py
test_add (__main__.TestArithmeticOperations.test_add) ... ok
test_subtract (__main__.TestArithmeticOperations.test_subtract) ... ok
test_multiply (__main__.TestArithmeticOperations.test_multiply) ... ok

Creating Test Fixtures

Great question! Let’s break this down into two parts:

🔹 `setUp()`, `tearDown()`, `setUpClass()`, `tearDownClass()` in `unittest`

These are test fixture methods used to prepare and clean up the environment for your tests.

1. `setUp()`

Runs before each test method.
Used to set up objects or state needed for the test.

2. `tearDown()`

Runs after each test method.
Used to clean up resources (e.g., closing files, stopping mocks).

3. `setUpClass()`

Runs once before all tests in the class.
Must be decorated with @classmethod.

4. `tearDownClass()`

Runs once after all tests in the class.
Also decorated with @classmethod.

✅ Example:

import unittest

class MyTestCase(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        print("🔧 setUpClass: Run once before all tests")

    @classmethod
    def tearDownClass(cls):
        print("🧹 tearDownClass: Run once after all tests")

    def setUp(self):
        print("🔧 setUp: Run before each test")

    def tearDown(self):
        print("🧹 tearDown: Run after each test")

    def test_one(self):
        print("✅ Running test_one")

    def test_two(self):
        print("✅ Running test_two")

Output:

🔧 setUpClass: Run once before all tests
🔧 setUp: Run before each test
✅ Running test_one
🧹 tearDown: Run after each test
🔧 setUp: Run before each test
✅ Running test_two
🧹 tearDown: Run after each test
🧹 tearDownClass: Run once after all tests

🔹 What is `MagicMock`?

MagicMock is a subclass of Mock from unittest.mock. It behaves like a mock object but also includes default implementations for magic methods (like __len__, __getitem__, __iter__, etc.).

✅ Use Case:

from unittest.mock import MagicMock

mock_obj = MagicMock()
mock_obj.__len__.return_value = 5

print(len(mock_obj))  # Output: 5

🔸 Why Use `MagicMock`?

When you’re mocking objects that use magic methods (dunder methods).
It saves you from manually defining those methods.

In Python unit testing, test fixtures are used to set up the environment before each test and clean it up afterward. When it comes to mocking instance methods and class methods, the unittest.mock module provides powerful tools to do this.

Here’s how you can mock both:

🔹 Mocking an Instance Method

Suppose you have a class like this:

class MyClass:
    def greet(self):
        return "Hello"

You can mock the greet method like this:

from unittest import TestCase
from unittest.mock import patch

class TestMyClass(TestCase):
    @patch.object(MyClass, 'greet', return_value='Mocked Hello')
    def test_greet(self, mock_greet):
        obj = MyClass()
        result = obj.greet()
        self.assertEqual(result, 'Mocked Hello')
        mock_greet.assert_called_once()

🔹 Mocking a Class Method

class MyClass:
    @classmethod
    def get_name(cls):
        return "Original Name"

You can mock it similarly:

from unittest import TestCase
from unittest.mock import patch

class TestMyClass(TestCase):
    @patch.object(MyClass, 'get_name', return_value='Mocked Name')
    def test_get_name(self, mock_get_name):
        result = MyClass.get_name()
        self.assertEqual(result, 'Mocked Name')
        mock_get_name.assert_called_once()

🔸 Key Points

Use @patch.object(ClassName, 'method_name') to mock methods.
For instance methods, you usually call them on an object.
For class methods, you call them on the class itself.
You can also use patch() as a context manager if you prefer not to use decorators.

✅ Bonus: Mocking with Fixtures

If you want to reuse mocks across multiple tests, you can use setUp():

class TestMyClass(TestCase):
    def setUp(self):
        patcher = patch.object(MyClass, 'greet', return_value='Hello from setup')
        self.mock_greet = patcher.start()
        self.addCleanup(patcher.stop)

    def test_greet(self):
        obj = MyClass()
        self.assertEqual(obj.greet(), 'Hello from setup')

Scenario Based Question

import datetime

def is_weekday():
    today = datetime.date.today()
    return 0 <= today.weekday() < 5
import weekday

class TestWeekday(unittest.TestCase):
    @patch("weekday.datetime")
    def test_is_weekday(self, mock_datetime):
        mock_datetime.date.today.return_value = datetime.date(2024, 4, 4)
        self.assertTrue(weekday.is_weekday())

    @patch("weekday.datetime")
    def test_is_weekend(self, mock_datetime):
        mock_datetime.date.today.return_value = datetime.date(2024, 4, 6)
        self.assertFalse(weekday.is_weekday())

Great questions! Let’s break down both parts clearly:

✅ 1. Why are we mocking `datetime`?

You’re mocking datetime.date.today() because:

The function is_weekday() depends on the current date.
In unit testing, you want predictable, repeatable results.
If you don’t mock it, the test result will change depending on the actual day it’s run.

🔍 Without mocking:

def test_is_weekday():
    assert is_weekday()  # Might pass or fail depending on today's date

✅ With mocking:

@patch("weekday.datetime")
def test_is_weekday(self, mock_datetime):
    mock_datetime.date.today.return_value = datetime.date(2024, 4, 4)  # Thursday
    self.assertTrue(weekday.is_weekday())

This ensures the test always runs with a known date (April 4, 2024), which is a Thursday (weekday).

✅ 2. Who attaches the second parameter `mock_datetime` to the test method?

That’s done by the @patch decorator from unittest.mock.

How it works:

@patch("weekday.datetime") replaces datetime in the weekday module with a mock object.
It automatically injects that mock object as the first argument to your test method.
So mock_datetime is the mock version of datetime inside the weekday module.

Behind the scenes:

@patch("weekday.datetime")
def test_is_weekday(self, mock_datetime):  # mock_datetime is injected here

You can name it anything (mock_dt, mocked_datetime, etc.) — it’s just a parameter that receives the mock.

✅ Why `@patch("weekday.datetime")` instead of `@patch("datetime.datetime")`?

🔹 The key idea: Patch where it’s used, not where it’s defined.

When you write:

@patch("weekday.datetime")

You’re patching the datetime as imported and used inside the weekday module — not the global datetime module.

🔍 Why this matters:

If you patch datetime.datetime globally, it affects all uses of datetime, which can lead to unintended side effects across your tests or even other modules.

But by patching weekday.datetime, you’re only mocking the datetime object as seen by the weekday module — keeping the scope clean and safe.

✅ Alternative: Mock the whole function

You’re absolutely right — another clean and often better approach is to mock the function directly, especially if it’s small and self-contained.

Example:

Instead of mocking datetime, you could do:

@patch("weekday.is_weekday", return_value=True)
def test_is_weekday(self, mock_is_weekday):
    self.assertTrue(weekday.is_weekday())

This avoids mocking internals and focuses on behavior — which is often more robust and readable.

W.I.P

Reference:-

Source link

Author Info

Find Me On

Trending News