Python Testing & Debugging Beyond unittestAdvanced Python testing playbooks
Isolation & Contracts

Patching an Async Context Manager

async with pool.acquire() as conn: is one of the most common lines in async Python and one of the easiest to mock incorrectly. The manager has two async protocol methods, the variable after as is bound to whatever __aenter__ returns rather than to the manager itself, and __aexit__'s return value decides whether an exception in the body propagates or vanishes. Get the last one wrong and the test passes while silently swallowing the very error it was written to observe.

These mistakes are especially costly because async context managers are almost always guarding a scarce resource — a pooled connection, a lock, a transaction, an open stream — and the whole point of the manager is that the resource is released on every exit path. A mock that suppresses exceptions or hands the body the wrong object makes a test pass precisely in the situation where the real resource would leak. The fix is a few explicit lines rather than reliance on auto-configured mocks, plus an assertion that cleanup actually ran. Once written as a small helper, the pattern is reusable for every pool, session, lock and client in an async codebase, and the helper itself becomes the place where the rule about __aexit__ is encoded once rather than remembered everywhere.

Prerequisites

Solution

A helper builds the manager correctly once; tests use it everywhere.

Python
from unittest.mock import AsyncMock, MagicMock

import pytest


def async_cm(value):
    """An async context manager double that behaves like a real one."""
    cm = MagicMock()
    cm.__aenter__ = AsyncMock(return_value=value)    # what `as` binds to
    cm.__aexit__ = AsyncMock(return_value=False)     # False: exceptions propagate
    return cm


async def test_failed_query_still_releases_the_connection(repo):
    connection = AsyncMock()
    connection.fetch.side_effect = RuntimeError("query failed")
    repo.pool.acquire = MagicMock(return_value=async_cm(connection))

    with pytest.raises(RuntimeError, match="query failed"):
        await repo.load_orders("cus_1")

    manager = repo.pool.acquire.return_value
    manager.__aexit__.assert_awaited_once()            # cleanup ran on the failure path
    exc_type, exc, _tb = manager.__aexit__.await_args.args
    assert exc_type is RuntimeError                    # and saw the real exception
The async with protocol, step by step Entering the block awaits __aenter__, whose return value is bound to the name after as. The body runs. On exit, __aexit__ is awaited with the exception type, value and traceback, or three Nones on success. If __aexit__ returns a truthy value the exception is suppressed; if it returns False the exception propagates to the caller. Two awaits and one decision await __aenter__() result bound by "as" body runs may raise await __aexit__(type, exc, tb) return value decides the exception's fate returns a Mock (truthy) exception suppressed — test passes wrongly returns False exception propagates — as in production
The left-hand outcome is the default for an unconfigured mock, which is why every async context manager double needs __aexit__ set explicitly.

Why this works

async with expands to an await on __aenter__, the body, and an await on __aexit__ with the exception triple — three None values on success. Python then checks the truthiness of what __aexit__ returned: truthy means "I handled it, suppress the exception", falsy means "let it propagate". A MagicMock's auto-created __aexit__ returns another mock object, which is truthy, so every exception raised in the body disappears.

Setting __aexit__ to an AsyncMock(return_value=False) restores the real manager's behaviour for any manager that does not deliberately suppress exceptions — which is nearly all of them. Recording the call also gives the test something to assert on: that cleanup was awaited, and with which exception, which is the evidence that the resource would have been released in production.

The binding of as deserves the same care. Python evaluates the expression after async with, awaits its __aenter__, and binds the result — so async with pool.acquire() as conn gives conn whatever __aenter__ returned, not the manager. A double that returns itself from __aenter__ works for managers like locks, where the manager and the resource are the same object, and silently misleads for pools and sessions, where the body expects a different object entirely. Configuring __aenter__'s return value explicitly, every time, removes the ambiguity.

Edge cases and failure modes

  • Forgetting __aenter__'s return value. The body receives a generic mock instead of the connection it expects, and assertions on it test nothing. Always set it explicitly.
  • Managers that should suppress. A few managers — contextlib.suppress, some retry helpers — legitimately return True. Mirror that deliberately and name it in the test.
  • acquire() being awaited itself. Some APIs are async with await pool.acquire() rather than async with pool.acquire(). The first needs acquire to be an AsyncMock returning the manager; the second needs a plain MagicMock. Match the real API.
  • Nested managers. A transaction inside a connection inside a pool needs each level configured. The helper composes: async_cm(async_cm(transaction)) is not right, but returning async_cm(conn) from acquire and async_cm(tx) from conn.transaction is.
  • Using a real manager when available. For your own managers, a fake implementation with real __aenter__/__aexit__ methods is often clearer than a mock and cannot get the protocol wrong.

Testing both exit paths

A context manager exists to guarantee cleanup, and the guarantee has two halves. The success path — body completes, __aexit__ receives three Nones, resources are released — is what every test exercises by accident. The failure path — body raises, __aexit__ receives the exception, resources are still released, the exception still reaches the caller — is what production exercises during an incident and tests exercise almost never.

Writing both as explicit tests turns the guarantee into something the suite checks rather than something the code claims. The failure-path test is the more valuable of the two, and it should assert three things: the exception propagated to the caller, __aexit__ was awaited exactly once, and it received the real exception type. That last check catches a subtle bug where intermediate code catches the original exception and raises a different one, leaving the manager to clean up with the wrong information — harmless for a connection, but significant for a transaction manager that decides between commit and rollback based on whether an exception occurred.

A third test is worth adding for any manager wrapping a transaction: that on the success path the commit happened and on the failure path the rollback did. With a hand-configured double this is a matter of giving the transaction mock commit and rollback as AsyncMocks and asserting which one was awaited. It is the test that proves the manager does its job rather than merely that Python called its methods, and it is the one that fails when someone refactors the manager and gets the condition backwards.

Assertions for the success and failure exits Two columns of assertions. On the success path, __aexit__ is awaited with three Nones and the transaction's commit is awaited. On the failure path, the exception reaches the caller, __aexit__ is awaited with the real exception type, and the transaction's rollback is awaited instead of commit. Two exits, two sets of assertions body succeeds __aexit__ awaited with (None, None, None) transaction.commit awaited rollback not awaited the path every test hits by accident body raises exception reaches the caller __aexit__ saw the real exception type rollback awaited, commit not the path production takes in an incident
The right-hand column is the one to write first. It covers the behaviour that matters most and is exercised least.

A fake manager instead of a mock

For managers you own, a small fake class is often clearer than any mock configuration, because it implements the protocol in ordinary Python and cannot get the suppression rule wrong by accident.

Python
class FakeConnection:
    def __init__(self):
        self.executed: list[str] = []
        self.released = False
        self.exit_exception: type | None = None

    async def execute(self, sql: str) -> None:
        self.executed.append(sql)


class FakePool:
    def __init__(self):
        self.connection = FakeConnection()

    def acquire(self):
        return self                              # the pool is its own manager here

    async def __aenter__(self):
        return self.connection

    async def __aexit__(self, exc_type, exc, tb):
        self.connection.released = True
        self.connection.exit_exception = exc_type
        return False                             # never suppress

Tests then read naturally — assert pool.connection.released, assert pool.connection.exit_exception is RuntimeError — with no knowledge of how mocks record awaits. The fake is also reusable across every test that needs a pool, which is where it pays back over repeated mock configuration, and it can be checked against the real pool with a small contract suite in the same way as any other fake.

The trade-off is a few dozen lines of test support code. For a manager used in one or two tests, the helper function above is enough. For the connection pool, the database session or the HTTP client that half the suite depends on, the fake is the better long-term choice — it is written once, reviewed once, and every test that uses it inherits a correct protocol implementation without having to know what one looks like, for the same reasons argued in writing an in-memory fake repository.

Mock helper versus fake manager Two approaches. The mock helper configures __aenter__ and __aexit__ on a MagicMock and suits managers used in a few tests. The fake manager implements the protocol in a small class that records release and the exit exception as plain attributes, suiting managers used widely across the suite. Choose by how many tests need the manager async_cm(value) helper a few lines, no new class asserts via await_args a manager used in a few tests FakePool class plain attributes: released, exit_exception, executed a manager half the suite uses
Both get the protocol right. The fake additionally makes the assertions read as statements about the resource rather than about mock bookkeeping.

Frequently Asked Questions

Why does my test pass even though the body raised inside async with? Because the mocked __aexit__ returned a truthy value. An unconfigured MagicMock or AsyncMock returns a Mock object, which is truthy, and a truthy return from __aexit__ tells Python to suppress the exception. Set __aexit__ to an AsyncMock with return_value=False so exceptions propagate as they would with the real manager.

What does async with bind the target variable to? The value returned by awaiting __aenter__, not the manager object itself. If a test configures the manager but forgets __aenter__'s return value, the body receives a generic Mock rather than the connection or session it expects, and assertions on that object test nothing meaningful.

Can autospec handle async context managers? Yes. create_autospec on a class that defines __aenter__ and __aexit__ produces a mock with AsyncMock versions of both, with the real signatures. You still need to set __aenter__'s return value, but __aexit__'s signature is enforced and the protocol is correctly async.

← Back to Patching Async Code & Coroutines