Context Managers and Testing

Introduction

Two habits separate scripts from professional code. The first is cleaning up reliably: files closed, locks released, connections ended, even when something goes wrong. The second is testing: proving that your code does what you claim, so you can change it without fear.

The with statement

You have used with open(...). It is an example of a context manager: an object that sets something up when the block starts and cleans up when it ends, whatever happens in between.

• __enter__ runs at the start. Its result is available after as.
• __exit__ runs at the end, always, and receives details about any exception. Returning True would swallow it.

contextlib: the shortcut

Writing a class is verbose. contextlib.contextmanager turns a generator into a context manager. Everything before yield is setup, and everything after is cleanup.

Use try/finally around the yield, so the cleanup runs even when the block raises.

A realistic example: temporarily change a setting, and always restore it.

Testing with assert

An assert statement states something that must be true. If not, it raises AssertionError.

Change one of the asserts so it is wrong and run it again to see the failure.

unittest

For larger programs, use a test framework. Python includes unittest. Write a class whose methods start with test_.

Each test_ method is a small experiment. The report shows how many passed. Change an expected value to see a failure with a helpful message. Useful assertions: assertEqual, assertTrue, assertIn and assertRaises.

Good test habits

• Test behaviour, not implementation: what goes in, what comes out.
• Cover the normal case, the edges (empty input, zero, the maximum) and the errors.
• Keep each test small and independent, with a name that says what it checks.
• When you find a bug, first write a test that fails because of it. After fixing it, the test stays and guards against the bug returning.
• In real projects many people use pytest, a popular third-party tool where tests are plain functions with assert. The ideas are identical.

Try it yourself

Write a function is_valid_password(text) that returns True when the text has at least 8 characters and contains at least one digit. Write unittest tests for a good password, a short one and one with no digit, then run them.

Summary

• A context manager does setup and guaranteed cleanup around a with block, using __enter__ and __exit__.
• @contextmanager builds one from a generator: setup, yield, then cleanup in finally.
• assert checks a claim, and unittest groups tests in test_ methods with helpful assertions.
• Test behaviour, cover edges and errors, and turn each bug into a test.