Context manager testing and debugging techniques
Testing context managers thoroughly ensures they handle both happy paths and edge cases: normal exit, exceptions, premature exit, and cleanup failures. Debugging context managers reveals subtle issues: resources not released, exceptions swallowed, or __exit__ order problems. This article covers practical testing strategies, common pitfalls, and tools for validating context manager behavior.
Proper testing of context managers is not optional — it is essential for writing reliable, maintainable code. A leaky context manager can silently corrupt application state, exhaust system resources, or cause cascading failures in long-running services.
Testing Strategy 1: Direct Testing with Context Manager Methods
The simplest approach is to call __enter__ and __exit__ directly, verifying their behavior without relying on the with statement.
# Example 1: Direct unit test of context manager
import unittest
from contextlib import contextmanager
class FileBackup:
"""A context manager that creates a backup before opening a file."""
def __init__(self, filepath):
self.filepath = filepath
self.backup_path = filepath + ".backup"
self.file_handle = None
def __enter__(self):
import shutil
import os
if os.path.exists(self.filepath):
shutil.copy(self.filepath, self.backup_path)
self.file_handle = open(self.filepath, 'w')
return self.file_handle
def __exit__(self, exc_type, exc_val, exc_tb):
if self.file_handle:
self.file_handle.close()
if exc_type is not None:
# Restore from backup on error
import shutil
import os
if os.path.exists(self.backup_path):
shutil.copy(self.backup_path, self.filepath)
print(f"Restored {self.filepath} from backup")
return False
class TestFileBackup(unittest.TestCase):
def setUp(self):
import tempfile
self.tmpdir = tempfile.mkdtemp()
self.test_file = f"{self.tmpdir}/test.txt"
def tearDown(self):
import shutil
shutil.rmtree(self.tmpdir)
def test_normal_usage(self):
"""Test successful file write with backup creation."""
# Setup: create initial file
with open(self.test_file, 'w') as f:
f.write("Original content")
# Use the context manager
cm = FileBackup(self.test_file)
file_handle = cm.__enter__()
file_handle.write("New content")
cm.__exit__(None, None, None)
# Verify
with open(self.test_file, 'r') as f:
self.assertEqual(f.read(), "New content")
import os
self.assertTrue(os.path.exists(f"{self.test_file}.backup"))
def test_exception_rollback(self):
"""Test that backup is restored on exception."""
# Setup
with open(self.test_file, 'w') as f:
f.write("Original content")
# Use context manager and raise exception
cm = FileBackup(self.test_file)
file_handle = cm.__enter__()
file_handle.write("Partial content")
cm.__exit__(ValueError, ValueError("test"), None)
# Verify original is restored
with open(self.test_file, 'r') as f:
self.assertEqual(f.read(), "Original content")
# Run the test
if __name__ == '__main__':
unittest.main(verbosity=2)
Testing Strategy 2: Using with Statement and assert_raises
Test the context manager using the with statement and verify exception handling.
# Example 2: Testing exception handling with with statement
import unittest
from contextlib import contextmanager
@contextmanager
def must_succeed():
"""A context manager that fails if any exception occurs."""
try:
yield
except Exception as e:
raise AssertionError(f"Unexpected error: {e}") from e
class TestContextManagerExceptions(unittest.TestCase):
def test_exception_is_reraised(self):
"""Verify exceptions are propagated."""
with self.assertRaises(ValueError):
with must_succeed():
raise ValueError("Test error")
def test_no_exception_succeeds(self):
"""Verify normal execution succeeds."""
with must_succeed():
result = 1 + 1
self.assertEqual(result, 2)
# Run
if __name__ == '__main__':
unittest.main(verbosity=2)
Testing Strategy 3: Mock and Spy Patterns
Use mocks to verify that cleanup methods are called and that resources are released.
# Example 3: Using mocks to verify cleanup
import unittest
from unittest.mock import Mock, patch, MagicMock
class DatabaseConnection:
"""A context manager for database connections."""
def __init__(self):
self.connection = None
def __enter__(self):
self.connection = Mock() # Simulate a connection
self.connection.connect()
return self.connection
def __exit__(self, exc_type, exc_val, exc_tb):
if self.connection:
self.connection.commit() if exc_type is None else self.connection.rollback()
self.connection.close()
return False
class TestDatabaseConnectionCleanup(unittest.TestCase):
def test_commit_on_success(self):
"""Verify commit is called on successful exit."""
cm = DatabaseConnection()
with patch.object(cm, '__enter__', return_value=Mock()):
with cm as mock_conn:
pass
# Alternatively, test directly
cm = DatabaseConnection()
conn_mock = cm.__enter__()
cm.__exit__(None, None, None)
# Verify methods were called
conn_mock.connect.assert_called_once()
conn_mock.commit.assert_called_once()
conn_mock.close.assert_called_once()
def test_rollback_on_exception(self):
"""Verify rollback is called on exception."""
cm = DatabaseConnection()
conn_mock = cm.__enter__()
cm.__exit__(ValueError, ValueError("test"), None)
# Verify rollback and close were called
conn_mock.rollback.assert_called_once()
conn_mock.close.assert_called_once()
conn_mock.commit.assert_not_called()
# Run
if __name__ == '__main__':
unittest.main(verbosity=2)
Testing Strategy 4: Pytest Fixtures
Pytest fixtures provide a natural way to test context managers.
# Example 4: Pytest fixtures for context manager testing
import pytest
import tempfile
import os
@pytest.fixture
def temp_file():
"""Fixture that provides a temporary file."""
fd, path = tempfile.mkstemp()
os.close(fd)
yield path
os.unlink(path)
class FileLocker:
"""A context manager that locks a file during use."""
def __init__(self, filepath):
self.filepath = filepath
self.locked = False
def __enter__(self):
with open(self.filepath, 'a') as f:
f.write("LOCKED\n")
self.locked = True
return self
def __exit__(self, exc_type, exc_val, exc_tb):
with open(self.filepath, 'a') as f:
f.write("UNLOCKED\n")
self.locked = False
return False
def test_file_locker_normal_exit(temp_file):
"""Test file locker with normal exit."""
with FileLocker(temp_file) as locker:
assert locker.locked
with open(temp_file, 'r') as f:
contents = f.read()
assert "LOCKED" in contents
assert "UNLOCKED" in contents
def test_file_locker_exception_exit(temp_file):
"""Test file locker cleanup on exception."""
try:
with FileLocker(temp_file) as locker:
raise ValueError("Test error")
except ValueError:
pass
with open(temp_file, 'r') as f:
contents = f.read()
# Even on exception, cleanup should run
assert "LOCKED" in contents
assert "UNLOCKED" in contents
# Run with: pytest example4.py -v
Testing Strategy 5: Parameterized Tests for Edge Cases
Test multiple scenarios with parameterized tests.
# Example 5: Parameterized testing
import pytest
from contextlib import contextmanager
@contextmanager
def exception_handler(suppress_types=None):
"""A context manager that selectively suppresses exceptions."""
suppress_types = suppress_types or ()
try:
yield
except suppress_types:
print("Suppressed")
@pytest.mark.parametrize("exception_type,should_suppress", [
(ValueError, True),
(TypeError, True),
(RuntimeError, False),
])
def test_exception_suppression(exception_type, should_suppress):
"""Test selective exception suppression."""
suppress_types = (ValueError, TypeError)
if should_suppress:
# Should not raise
with exception_handler(suppress_types):
raise exception_type("test")
else:
# Should raise
with pytest.raises(exception_type):
with exception_handler(suppress_types):
raise exception_type("test")
Common Pitfalls and Debugging
Pitfall 1: Forgetting to Clean Up in exit
# WRONG: cleanup not guaranteed
class BadResource:
def __enter__(self):
self.resource = open("file.txt")
return self.resource
# Missing __exit__!
# CORRECT: cleanup guaranteed
class GoodResource:
def __enter__(self):
self.resource = open("file.txt")
return self.resource
def __exit__(self, exc_type, exc_val, exc_tb):
if self.resource:
self.resource.close()
return False
Pitfall 2: Suppressing Exceptions Silently
# WRONG: silently suppresses all exceptions
class BadContext:
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
return True # Suppresses!
# CORRECT: propagates exceptions
class GoodContext:
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
return False # Propagates
Pitfall 3: Not Handling Partial Acquisition
# WRONG: cleanup failure if acquisition fails partway
class BadAcquisition:
def __enter__(self):
self.resource_a = acquire_a()
self.resource_b = acquire_b() # Might fail!
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.resource_a.close()
self.resource_b.close()
# CORRECT: partial cleanup inside __enter__
class GoodAcquisition:
def __enter__(self):
try:
self.resource_a = acquire_a()
self.resource_b = acquire_b()
return self
except Exception:
if hasattr(self, 'resource_a'):
self.resource_a.close()
raise
def __exit__(self, exc_type, exc_val, exc_tb):
self.resource_a.close()
self.resource_b.close()
return False
Debugging Tools
Tool 1: Print Statements for Flow Tracking
# Simple but effective
class DebugContext:
def __enter__(self):
print("DEBUG: __enter__ called")
return self
def __exit__(self, exc_type, exc_val, exc_tb):
print(f"DEBUG: __exit__ called with exc_type={exc_type}")
return False
Tool 2: Using pdb Breakpoints
import pdb
class DebuggableContext:
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is not None:
pdb.post_mortem(exc_tb) # Drop into debugger on exception
return False
Tool 3: Logging
import logging
class LoggedContext:
def __init__(self, name):
self.name = name
self.logger = logging.getLogger(name)
def __enter__(self):
self.logger.info(f"{self.name}: acquiring")
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type:
self.logger.error(f"{self.name}: error {exc_type.__name__}")
else:
self.logger.info(f"{self.name}: cleanup")
return False
Key Takeaways
- Test context managers by calling
__enter__and__exit__directly for fine-grained control. - Use
withstatement in tests to verify behavior in realistic scenarios. - Mock cleanup methods to verify they are called with the correct arguments.
- Test both normal exit and exception paths for each context manager.
- Use parameterized tests to cover edge cases and multiple exception types.
- Always ensure cleanup runs, even on exceptions and partial acquisition failures.
Frequently Asked Questions
How do I test if a resource was leaked?
Use mocks to track cleanup calls. Verify that __exit__ methods are called and that all acquired resources call their release methods. Use resource monitoring tools like psutil to check system resources.
Should I test enter and exit separately or with the with statement?
Both. Test them separately for fine-grained verification, and test with the with statement for integration testing. This covers both isolated and realistic scenarios.
How do I debug context manager cleanup order?
Add logging to each context manager's __exit__ method. Use a common logger with timestamps to see the exact order of cleanup. Stack traces in exceptions also show the nesting order.
Can I use context managers in fixtures?
Yes. Use yield in a pytest fixture, which is itself a context manager. The code after yield is the cleanup, analogous to __exit__.
How do I test async context managers?
Use pytest-asyncio or similar async testing framework. Mark your test async def and use async with in tests. Verify cleanup with mocks and assertions on async methods.