Python Unittest Example

Python's `unittest` library provides a framework for writing and running tests. It allows developers to ensure code works as expected, helping catch bugs early in development. In this guide, we’ll cover the basics and advanced features of `unittest` with multiple scenarios.

1. Setting Up a Basic Unit Test

A unit test checks small parts of code, like individual functions, to ensure they return correct results. Begin by importing `unittest` and defining a test class that inherits from `unittest.TestCase`.

Example: Testing a function that adds two numbers.
# mymodule.py
def add(x, y):
    """Returns the sum of x and y."""
    return x + y
# test_mymodule.py
import unittest
from mymodule import add

class TestAddFunction(unittest.TestCase):
    
    def test_add_positive_numbers(self):
        """Test addition with positive integers."""
        self.assertEqual(add(2, 3), 5)
    
    def test_add_negative_numbers(self):
        """Test addition with negative integers."""
        self.assertEqual(add(-1, -1), -2)
    
    def test_add_zero(self):
        """Test addition with zero."""
        self.assertEqual(add(0, 0), 0)

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

Output:
Ran 3 tests in 0.001s
OK

Explanation:
Each `test_*` method in `TestAddFunction` runs as an individual test case. `self.assertEqual` checks if the actual output matches the expected output. All tests pass, meaning `add()` works as expected for each input scenario.

2. Testing Exceptions

You can test whether code raises an expected exception using `assertRaises`.

Example: Testing an exception raised when dividing by zero.
# mymodule.py
def divide(x, y):
    """Divide x by y."""
    if y == 0:
        raise ValueError("Cannot divide by zero")
    return x / y
# test_mymodule.py
import unittest
from mymodule import divide

class TestDivideFunction(unittest.TestCase):
    
    def test_divide_by_zero(self):
        """Test division by zero raises ValueError."""
        with self.assertRaises(ValueError):
            divide(10, 0)

    def test_normal_division(self):
        """Test normal division."""
        self.assertEqual(divide(10, 2), 5)

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

Output:
Ran 2 tests in 0.001s
OK

Explanation:
In `test_divide_by_zero`, the `assertRaises` context checks if `ValueError` is raised. If `divide(10, 0)` raises the expected exception, the test passes. `test_normal_division` confirms `divide()` works normally with non-zero divisors.

3. Using `setUp` and `tearDown` Methods

`setUp` and `tearDown` methods are called before and after each test method, respectively. These methods are useful for initializing objects or cleaning up resources.

Example: Using `setUp` and `tearDown` with a file.
import unittest

class TestFileOperations(unittest.TestCase):

    def setUp(self):
        """Create a file for testing."""
        self.file = open("testfile.txt", "w")
        self.file.write("Hello, unittest!")
        self.file.close()

    def test_read_file(self):
        """Test reading from the file."""
        with open("testfile.txt", "r") as f:
            content = f.read()
        self.assertEqual(content, "Hello, unittest!")

    def tearDown(self):
        """Remove the file after testing."""
        import os
        os.remove("testfile.txt")

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

Output:
Ran 1 test in 0.001s
OK

Explanation:
`setUp` creates and writes to `testfile.txt` before each test, while `tearDown` deletes it after. This ensures a fresh test environment each time.

4. Parameterized Tests with `subTest`

`subTest` allows multiple related tests within a single test function, which is useful for testing the same logic with different parameters.

Example: Testing multiple cases in a single test function.
import unittest
from mymodule import add

class TestAddFunction(unittest.TestCase):

    def test_add_multiple_cases(self):
        """Test addition with multiple cases."""
        cases = [
            (2, 3, 5),
            (-1, -1, -2),
            (0, 0, 0),
            (100, 200, 300)
        ]
        for x, y, expected in cases:
            with self.subTest(x=x, y=y, expected=expected):
                self.assertEqual(add(x, y), expected)

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

Output:
Ran 1 test in 0.001s
OK

Explanation:
Each `subTest` tests `add()` with different parameters. If a test fails, only that specific `subTest` fails, while others continue to run.

5. Mocking in `unittest`

Mocking allows you to replace parts of your system under test, such as functions or objects, with mock versions to test code in isolation.

Example: Mocking a function to test without real database access.
# mymodule.py
import requests

def get_status(url):
    """Fetch status code for a given URL."""
    response = requests.get(url)
    return response.status_code
# test_mymodule.py
import unittest
from unittest.mock import patch
from mymodule import get_status

class TestGetStatus(unittest.TestCase):
    
    @patch("mymodule.requests.get")
    def test_get_status(self, mock_get):
        """Mock the requests.get call."""
        mock_get.return_value.status_code = 200
        self.assertEqual(get_status("http://example.com"), 200)

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

Output:
Ran 1 test in 0.001s
OK

Explanation:
The `@patch` decorator replaces `requests.get` with a mock object, allowing us to set `status_code` to 200 without a real HTTP request. This test confirms `get_status()` handles the expected response.

6. Running Tests with `unittest`

There are multiple ways to run tests with `unittest`:
- Running tests by script: Run `python test_mymodule.py`.
- Using the command line: Use `python -m unittest discover` to discover and run all test files.
- Verbose mode: Use `python -m unittest -v` for detailed output of each test.

Output:
test_add_positive_numbers (test_mymodule.TestAddFunction) ... ok
test_add_zero (test_mymodule.TestAddFunction) ... ok
test_add_negative_numbers (test_mymodule.TestAddFunction) ... ok
Ran 3 tests in 0.001s
OK

Explanation:
Using `-v` provides detailed output for each test case, indicating which tests passed or failed.

Conclusion

Python’s `unittest` framework provides a powerful and flexible testing solution, covering test cases, assertions, exception testing, setup and teardown, parameterized tests, and mocking. Mastering these elements ensures you can rigorously test your code and catch issues early in development.

Previous: Unit Test Introduction | Next: Unit Test Mocking

<
>