Python Testing & Debugging Beyond unittestAdvanced Python testing playbooks
Async & Concurrency

Testing Async Generators and Context Managers

An async generator that streams rows and closes its cursor in a finally looks correct and is routinely broken: a caller that breaks out of the async for leaves the generator suspended, the finally does not run, and the cursor stays open until garbage collection — possibly after the loop has closed, which raises during interpreter shutdown. Tests for these objects have to drive the whole lifecycle, including the paths that end early.

Prerequisites

  • Python 3.10+ for contextlib.aclosing; 3.7+ for asynccontextmanager.
  • pytest >= 8.0 with a runner configured per pytest-asyncio in depth.
  • An understanding that cancellation and GeneratorExit both arrive as exceptions at the suspension point.

Solution

Test three lifecycles for a generator — full consumption, early exit, and failure inside the body — and assert on cleanup each time.

Python
import contextlib

import pytest


async def stream_rows(pool):
    """Yields rows, and must always release the connection."""
    connection = await pool.acquire()
    try:
        async for row in connection.cursor("SELECT id FROM widget"):
            yield row
    finally:
        await pool.release(connection)      # must run on every exit path


async def test_full_consumption_releases_the_connection(pool):
    rows = [row async for row in stream_rows(pool)]
    assert len(rows) == 3
    assert pool.in_use == 0


async def test_early_exit_releases_the_connection(pool):
    # aclosing() guarantees aclose() at the end of the block, which throws
    # GeneratorExit into the suspended generator and runs its finally.
    async with contextlib.aclosing(stream_rows(pool)) as stream:
        async for _row in stream:
            break                            # abandon after one row

    assert pool.in_use == 0                  # fails without aclosing()


async def test_failure_inside_the_body_still_releases(pool):
    with pytest.raises(ValueError):
        async with contextlib.aclosing(stream_rows(pool)) as stream:
            async for _row in stream:
                raise ValueError("consumer blew up")

    assert pool.in_use == 0
Three exit paths from an async generator Three columns. Full consumption runs the generator to exhaustion and its finally block executes. Early exit without aclosing leaves the generator suspended so cleanup is deferred to garbage collection. Early exit inside an aclosing block throws GeneratorExit into the generator immediately, so cleanup runs at the end of the block. Where the finally block runs, and when consumed fully loop runs to exhaustion StopAsyncIteration raised finally runs immediately connection released the path everyone tests break, no aclosing generator left suspended finally deferred to GC may run after loop close connection still held the path that leaks break, with aclosing aclose() at block exit GeneratorExit thrown in finally runs deterministically connection released the path to write
The middle column is the default behaviour of a plain async for with a break, which is why the leak is so common and so rarely noticed in tests.

Why this works

An async generator suspended at a yield has no way to know the consumer has stopped. aclose() is what tells it: the coroutine throws GeneratorExit in at the suspension point, the finally executes, and the generator is marked closed. contextlib.aclosing is a context manager whose only job is to call aclose() on exit, which makes the cleanup deterministic and tied to a lexical block rather than to the garbage collector.

Without it, cleanup happens when the generator object is finalised. CPython will attempt that through the loop's asynchronous-generator finalisation hooks, but only while the loop is still running — after the loop closes, the pending finaliser produces RuntimeError: Event loop is closed during shutdown, or simply never runs. Which of those you get depends on timing, which is why the symptom is intermittent.

Edge cases and failure modes

  • Awaiting inside the finally after cancellation. A cleanup that awaits while the generator is being cancelled is itself cancelled. Wrap the essential part in asyncio.shield with its own deadline, as in timeouts, cancellation and deadlines.
  • yield inside a try/finally inside a lock. If the consumer abandons the generator, the lock is held until aclose(). This is the same bug as the connection leak with worse consequences.
  • Reusing an exhausted generator. A second async for over the same object yields nothing rather than restarting. Tests that reuse a generator fixture across two tests silently get an empty stream in the second.
  • __aexit__ returning a truthy mock. In a test that replaces the manager with a mock, an unconfigured __aexit__ returns a Mock, which is truthy and therefore suppresses the exception — see patching an async context manager.
  • asynccontextmanager over a generator that yields twice. It raises RuntimeError: generator didn't stop, which is confusing but literal: the manager expects exactly one yield.

Testing an async context manager's failure path

@asynccontextmanager turns a generator into a manager, and the same lifecycle questions apply with one addition: what happens to an exception raised in the body.

Python
import contextlib

import pytest


@contextlib.asynccontextmanager
async def transaction(connection):
    tx = connection.transaction()
    await tx.start()
    try:
        yield tx
    except Exception:
        await tx.rollback()          # the path most tests never exercise
        raise                        # re-raise: suppression would hide the bug
    else:
        await tx.commit()


async def test_body_failure_rolls_back(connection):
    with pytest.raises(ValueError):
        async with transaction(connection) as tx:
            await connection.execute("INSERT INTO widget (name) VALUES ('x')")
            raise ValueError("business rule violated")

    # The assertion that matters: state, not that rollback was called.
    assert await connection.fetchval("SELECT count(*) FROM widget") == 0


async def test_success_commits(connection):
    async with transaction(connection):
        await connection.execute("INSERT INTO widget (name) VALUES ('y')")

    assert await connection.fetchval("SELECT count(*) FROM widget") == 1

The raise after rollback() is easy to omit and catastrophic when omitted: the generator would swallow the exception, the caller would proceed as if the operation succeeded, and the test above is the only thing that would catch it. That asymmetry — one missing keyword converting an error into silent data loss — is why the failure path deserves a test of its own rather than being assumed.

Success and failure paths through an async context manager A single manager with two paths. On success, the body completes, the else branch commits and control returns normally. On failure, the exception enters the except branch, the transaction is rolled back and the exception is re-raised so the caller sees it. A note marks that omitting the re-raise silently suppresses the error. Two exits, one of which is usually untested yield tx — body runs no exception raises else: await tx.commit() except: await tx.rollback() raise — omit this and it is swallowed caller continues normally
Both branches need a test. The right-hand one is where a missing raise turns a failed operation into a silently committed one.

Asserting on what the stream produced

Generators tempt tests into checking only the values that came out, which misses two properties worth holding.

The first is laziness. A generator that eagerly builds the whole result before yielding anything defeats the purpose of streaming, and the test for it is a counter on the source:

Python
async def test_stream_is_lazy(pool, instrumented_cursor):
    async with contextlib.aclosing(stream_rows(pool)) as stream:
        first = await anext(stream)

    assert first is not None
    # Only one batch should have been fetched, not the whole table.
    assert instrumented_cursor.fetch_calls == 1

The second is ordering and completeness together. Collecting into a list and comparing to an expected list covers both in one assertion, and is preferable to checking length and membership separately — a stream that duplicates one row and drops another passes both of those and fails the list comparison.

What a lazy stream does differently from an eager one Two rows. The lazy generator fetches one batch, yields its rows, and fetches the next batch only when the consumer asks, so memory stays flat and the first row arrives immediately. The eager implementation fetches every batch before yielding anything, so the first row is delayed and memory grows with the result size. Laziness is a property a test can assert lazy fetch batch 1 → yield rows → consumer asks → fetch batch 2 → … first row immediately · memory flat · fetch_calls == 1 after one row eager fetch every batch → build a list → yield from it first row delayed · memory grows with the result · fetch_calls == N
Both implementations yield identical values, so a test that only compares outputs cannot tell them apart — and the difference is the entire reason the generator exists.

Where the stream is consumed by a pipeline rather than by a list comprehension, the useful assertion moves downstream: assert the consumer processed items incrementally, for example by checking that a progress callback fired more than once before the stream ended. That is the behavioural statement of the same property, and it survives a refactor of the generator's internals in a way the fetch_calls counter does not.

One more assertion is worth adding wherever the stream can be long: that the generator stops when told to. A consumer that sets a stop flag, or a deadline that fires, should end the iteration promptly rather than after the current batch of ten thousand rows. Testing it means driving the generator with a small batch size, signalling a stop after the first item, and asserting the source was not queried again — which is the same shape as the laziness test with the condition inverted.

Fixtures that are themselves async generators

A pytest_asyncio.fixture written with yield is an async generator, and the same lifecycle rules apply to it — with the difference that pytest-asyncio handles the closing, correctly, as long as the loop scopes line up.

Python
import pytest_asyncio


@pytest_asyncio.fixture(loop_scope="module")
async def stream(pool):
    # The plugin calls aclose() during teardown, on the fixture's loop.
    generator = stream_rows(pool)
    try:
        yield generator
    finally:
        await generator.aclose()     # explicit: do not rely on finalisation

Closing explicitly in the fixture's own finally is worth the extra line. It makes the teardown ordering visible, it runs while the loop is guaranteed alive, and it does not depend on the plugin's finalisation behaviour staying the same across versions.

The failure this avoids is a familiar one: a fixture whose generator is finalised after its loop has closed produces an exception during teardown that pytest reports against the next test, or as an error with no test attached at all. Making the close explicit and scope-matched keeps the teardown inside the window where it can succeed, which is the same principle that governs pools and servers in sharing an event loop across a test module.

Frequently Asked Questions

Why does my async generator's finally block not run? Because the generator was abandoned rather than closed. Breaking out of an async for leaves the generator suspended, and its cleanup runs only when aclose() is called or the object is finalised — which may be much later, on a different loop, or never. Use contextlib.aclosing() so the generator is closed deterministically at the end of the block.

How do I test that an async context manager cleans up when the body raises? Raise deliberately inside the block, catch the exception outside it, and then assert on the side effects cleanup was supposed to produce — a released connection, a rolled-back transaction, a closed file. Asserting that __aexit__ was called proves only that Python ran it, not that it did the right thing.

Should aexit ever return True? Only when the manager genuinely exists to swallow a specific exception, which is rare. A truthy return suppresses the exception and hides failures from every caller, including tests. Return False, or nothing at all, unless suppression is the manager's documented purpose.

← Back to pytest-asyncio in Depth