Write Simple Unit Tests

Have you ever fixed one bug only to discover you broke two other features? That sinking feeling is the reason software teams adopt automated testing. Without unit tests, every change is a gamble—you either manually re-test everything (and miss things), or you push code and hope for the best. The unittest module, built right into Python’s standard library, gives you a safety net: write small, fast tests once, and run them forever to catch regressions instantly.

The problem this lesson solves

Manual testing doesn’t scale. When your project grows past a single script, the number of possible interactions explodes. Think of it like juggling: you can keep three balls in the air by hand, but what about thirty? Even a tiny change—renaming a variable or adjusting a loop—can silently break a function that worked perfectly yesterday. Unit tests solve exactly this problem: they isolate a single unit of code (like a function or class method) and verify it produces the correct output for known inputs. This lesson walks you through writing those tests with Python’s built-in unittest framework, so you gain confidence that your code does what you expect, every time you run it.

Core concept / mental model

Picture a test harness like a workshop jig that holds a piece of metal steady while you drill a precise hole. The metal is your function, the drill bit is your test, and the jig is the unittest framework. A unit test has three parts:

1. Arrange – Set up the inputs and state.
2. Act – Call the code you’re testing.
3. Assert – Check that the result matches the expected value.

In unittest, each test is a method inside a class that inherits from unittest.TestCase. The class groups related tests; the methods run independently. When you call self.assertEqual() or self.assertTrue(), the framework records a pass or fail automatically.

Mental model shorthand: A unit test is a recipe. The ingredients are inputs; the steps are the function call; the tasting spoon is the assertion. If the dish tastes right, the test passes. If not, you get a clear error message telling you what went wrong.

How it works step by step

1. Import the module and create a test class

import unittest
from my_module import add  # assuming you have a function add

Every test class must inherit from unittest.TestCase. The class name can be anything, but following the pattern TestSomething is standard.

2. Write test methods starting with test_

class TestAddFunction(unittest.TestCase):
    def test_add_positive_numbers(self):
        result = add(2, 3)
        self.assertEqual(result, 5)

Any method starting with test_ is automatically discovered and run. Methods without that prefix are ignored by the test runner.

3. Choose the right assertion

4. Run your tests

There are two common ways: - From the command line: python -m unittest test_module.py - Directly in the script: add if __name__ == '__main__': unittest.main() at the bottom.

Hands-on walkthrough

Let’s say you have a module calculator.py with two functions:

# calculator.py
def add(x, y):
    return x + y

def divide(x, y):
    if y == 0:
        raise ValueError("Cannot divide by zero")
    return x / y

Now write unit tests in test_calculator.py:

import unittest
from calculator import add, divide

class TestCalculator(unittest.TestCase):

    def test_add_two_positive_numbers(self):
        self.assertEqual(add(10, 20), 30)

    def test_add_negative_and_positive(self):
        self.assertEqual(add(-5, 8), 3)

    def test_divide_normal_case(self):
        self.assertEqual(divide(10, 2), 5.0)

    def test_divide_by_zero_raises_error(self):
        with self.assertRaises(ValueError) as context:
            divide(10, 0)
        self.assertEqual(str(context.exception), "Cannot divide by zero")

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

Expected output when you run python -m unittest test_calculator.py

....
----------------------------------------------------------------------
Ran 4 tests in 0.001s

OK

Each . represents a passing test. If a test fails, you’d see an F and a traceback.

Pro tip: Use self.assertRaises as a context manager to test that exceptions are raised exactly where you expect them. That pattern is cleaner than a bare try/except block inside your test.

Compare options / when to choose what

While unittest is the built-in choice, Python developers often reach for other tools. Here’s a comparison:

When should you use unittest? When you want zero dependencies, are working in a large codebase that already uses it, or are submitting code to environments where only the standard library is allowed (e.g., many competitive programming platforms). When you want concise, readable tests with minimal boilerplate, pytest is the modern favorite. For this lesson, we stick with unittest because it teaches the fundamental patterns that all Python testing frameworks share.

Troubleshooting & edge cases

Common mistake: Forgetting to inherit from TestCase

class TestAdd:   # Wrong! Missing (unittest.TestCase)
    def test_add(self):
        assert add(1, 1) == 2

This class will be discovered but not run by the test runner, leaving you with zero tests executed—silent failure. Always check that your test class parent is unittest.TestCase.

Edge case: Floating-point comparisons

self.assertEqual(0.1 + 0.2, 0.3)  # May fail due to floating-point precision

Use assertAlmostEqual instead:

self.assertAlmostEqual(0.1 + 0.2, 0.3, places=7)

Gotcha: setUp vs setUpClass

setUp() runs before every test method—use it when each test needs a fresh copy of shared resources. setUpClass() runs once per class—use it for expensive setup like database connections. Be careful: shared mutable state between tests can cause mysterious failures.

Edge case: Testing code that calls exit()

A function that calls sys.exit() will kill the test runner. Test it this way:

with self.assertRaises(SystemExit):
    my_function_that_exits()

What you learned & what's next

You now understand the core idea behind unit testing with unittest: - Isolation: Each test checks one unit of code independently. - Automation: Once written, tests run in seconds and can be integrated into CI pipelines. - Feedback: Failures pinpoint exactly what broke—function, input, and expected vs. actual output.

You completed a practical exercise: testing an add and divide function, including the crucial edge case of division by zero. You also saw how to choose between unittest and pytest. In the next lesson, you’ll learn how to test more complex scenarios—functions that modify data structures or depend on external state—and you’ll apply the same unittest patterns to build a reliable test suite for a small project. Keep writing tests: they’re the best documentation your code will ever have.