Python
Context Manager
Mocking
Unit Testing
Python Testing

Python Mocking a context manager

Interview Questions practice on Codemia

Over 8,000 real interview questions from top companies, searchable by company and role.

Browse interview questions

Introduction

Context managers are common in Python because they wrap setup and cleanup in a readable with block. When you test code that opens files, database sessions, or network clients this way, the main challenge is not mocking the object itself, but mocking what __enter__() returns.

How Context Manager Mocking Works

When Python executes with something() as value:, it calls something().__enter__() and assigns that return value to value. That means your test usually needs to configure the mock one level deeper than people first expect.

Consider a simple function that reads a file:

python
def load_settings(path):
    with open(path, "r", encoding="utf-8") as handle:
        return handle.read().strip()

If you patch open, the object used inside the block is not the patch directly. It is the value returned by open(...).__enter__().

python
1from unittest.mock import patch, MagicMock
2
3def load_settings(path):
4    with open(path, "r", encoding="utf-8") as handle:
5        return handle.read().strip()
6
7def test_load_settings():
8    fake_open = MagicMock()
9    fake_open.return_value.__enter__.return_value.read.return_value = "debug=true\n"
10
11    with patch("builtins.open", fake_open):
12        result = load_settings("settings.txt")
13
14    assert result == "debug=true"

MagicMock is useful here because it already supports magic methods such as __enter__ and __exit__.

Prefer mock_open for File Access

For file-oriented tests, mock_open is usually clearer than building the chain by hand. It knows how file objects behave and keeps the test focused on the data.

python
1from unittest.mock import patch, mock_open
2
3def first_line(path):
4    with open(path, "r", encoding="utf-8") as handle:
5        return handle.readline().strip()
6
7def test_first_line():
8    mocked_file = mock_open(read_data="alpha\nbeta\n")
9
10    with patch("builtins.open", mocked_file):
11        result = first_line("demo.txt")
12
13    assert result == "alpha"

This is easier to maintain than configuring several nested attributes yourself. It also makes the intent obvious: the test cares about file content, not about the internals of open.

Mocking a Custom Context Manager

Many production codebases use custom context managers for database sessions, transactions, or API clients. In that case you often patch the class and configure its return_value.__enter__.return_value.

python
1from unittest.mock import patch, MagicMock
2
3class DatabaseSession:
4    def __enter__(self):
5        raise NotImplementedError
6
7    def __exit__(self, exc_type, exc, tb):
8        return False
9
10def fetch_usernames():
11    with DatabaseSession() as session:
12        return session.list_usernames()
13
14def test_fetch_usernames():
15    fake_session = MagicMock()
16    fake_session.list_usernames.return_value = ["alice", "bob"]
17
18    with patch("__main__.DatabaseSession") as mock_session_class:
19        mock_session_class.return_value.__enter__.return_value = fake_session
20        result = fetch_usernames()
21
22    assert result == ["alice", "bob"]

This pattern is the one to remember when the with target is a class instance rather than a file handle.

Verifying Cleanup Behavior

Sometimes the real behavior you want to test is cleanup. A context manager may commit, roll back, close a socket, or release a lock. In those cases, assert both the return value and the exit behavior.

python
1from unittest.mock import patch, MagicMock
2
3class Locker:
4    def __enter__(self):
5        raise NotImplementedError
6
7    def __exit__(self, exc_type, exc, tb):
8        return False
9
10def protected_read():
11    with Locker() as lock:
12        return lock.read()
13
14def test_protected_read():
15    fake_lock = MagicMock()
16    fake_lock.read.return_value = "ok"
17
18    with patch("__main__.Locker") as mock_locker:
19        instance = mock_locker.return_value
20        instance.__enter__.return_value = fake_lock
21
22        assert protected_read() == "ok"
23        instance.__enter__.assert_called_once()
24        instance.__exit__.assert_called_once()

That gives you confidence that the code actually used the context manager correctly, not just that it returned the expected data.

Common Pitfalls

The most common mistake is configuring mock.read.return_value and expecting the with block to use it. That fails because the code reads from mock.__enter__().read(), not from the top-level mock.

Another frequent issue is patching the wrong import path. Patch the name as it is used by the code under test, not necessarily where the object was originally defined. If a module imports open_client locally, patch that module’s reference.

Tests can also become brittle if they over-specify every call on the mock chain. Configure only the pieces your function actually depends on. That keeps refactors from breaking unrelated assertions.

Finally, do not reach for a plain Mock when magic methods are involved. Use MagicMock, mock_open, or a small fake object that implements __enter__ and __exit__.

Summary

  • A with block uses the value returned by __enter__(), so that is usually what your test must configure.
  • Use mock_open for file access because it is simpler and more readable than building nested mocks manually.
  • For custom context managers, patch the class and set return_value.__enter__.return_value.
  • Assert cleanup behavior when resource release matters, not just the final return value.
  • Patch the symbol where your code looks it up, or the mock will never be used.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

Interview Questions practice on Codemia

Over 8,000 real interview questions from top companies, searchable by company and role.

Browse interview questions

All Rights Reserved.