Testing Cancellation and Cleanup Paths

Cancellation is the path production takes whenever a deadline fires, a client disconnects or a shutdown begins, and it is the path test suites cover least. The usual state of affairs is a finally block that has never executed under cancellation, a lock that is released only on the happy path, and a connection that leaks once per timeout — discovered when the pool exhausts under load rather than in a test.

Prerequisites

• Python 3.11+ for asyncio.timeout; 3.8+ for CancelledError inheriting from BaseException.
• pytest >= 8.0 with an async runner configured per pytest-asyncio in depth.
• An understanding that cancellation is delivered as an exception at the next await, not as a thread kill.

Solution

Start the work as a task, cancel it at a point the test controls, and assert on the state cleanup was supposed to restore.

import asyncio

import pytest


class Worker:
    def __init__(self, pool, lock):
        self._pool, self._lock = pool, lock
        self.started = asyncio.Event()

    async def run(self):
        connection = await self._pool.acquire()
        await self._lock.acquire()
        try:
            self.started.set()            # the test's synchronisation point
            while True:
                await asyncio.sleep(3600) # a suspension point cancellation can reach
        finally:
            self._lock.release()          # must happen on the cancelled path too
            await self._pool.release(connection)


async def test_cancellation_releases_every_resource(pool, lock):
    worker = Worker(pool, lock)
    task = asyncio.create_task(worker.run())

    # Deterministic: wait for the worker's own signal, never a sleep.
    await asyncio.wait_for(worker.started.wait(), timeout=1)

    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task                        # awaiting lets the unwinding finish

    # The assertions that matter are about state, not about the exception.
    assert not lock.locked(), "the lock was not released on cancellation"
    assert pool.in_use == 0, "the connection was not returned to the pool"

Why this works

task.cancel() does not stop anything immediately. It arranges for CancelledError to be raised inside the coroutine at its next suspension point, which unwinds the stack and runs every finally on the way out. Because that unwinding is itself asynchronous, the test must await task before asserting — otherwise the assertions run while the finally may not have executed.

Waiting on the worker's own started event rather than sleeping is what makes the cancellation land at a known place. A sleep-based version cancels at whatever point the scheduler happens to be at, which means the test sometimes cancels before the resources are acquired and therefore proves nothing.

Edge cases and failure modes

• A coroutine blocked in synchronous code. time.sleep, a blocking socket read or a CPU loop never reaches an await, so the cancellation is never delivered. Push the work to asyncio.to_thread.
• except Exception around the await. It does not catch CancelledError in 3.8+, which is correct — but except BaseException and bare except: do, and both silently absorb cancellation.
• await inside finally without shielding. Once cancellation is in flight, a further await in the cleanup is cancelled immediately, so the cleanup never completes.
• Asserting before awaiting the task. The unwinding has not finished; the assertion is a race.
• Cancelling a task that already finished. cancel() returns False and nothing happens, so a test that cancels too late passes vacuously. Assert on task.cancelled() when the distinction matters.

Making cancellation reachable in the first place

Half the cancellation bugs in a codebase are not bugs in the cleanup at all — they are places where the cancellation can never be delivered, so the cleanup is irrelevant. Those are worth finding before writing any assertion about unwinding.

A coroutine is cancellable only at its suspension points. Between them it runs to completion regardless of what anyone requests. Three patterns remove every suspension point from a region and therefore make it uncancellable:

import asyncio
import time

import requests


async def uncancellable_three_ways(rows):
    # 1. Blocking I/O: the loop itself stops, so nothing can be delivered.
    response = requests.get("https://slow.example")     # no await anywhere

    # 2. A CPU loop: no await, so no suspension point for seconds at a time.
    digest = sum(hash(row) for row in rows)             # rows may be enormous

    # 3. A blocking sleep, which is both of the above at once.
    time.sleep(5)
    return response, digest

import asyncio


async def cancellable_version(rows, client):
    # 1. A real async client yields at every network operation.
    response = await client.get("https://slow.example")

    # 2. CPU work moved to a thread; the await is a suspension point, and the
    #    thread finishes on its own after cancellation rather than being killed.
    digest = await asyncio.to_thread(lambda: sum(hash(row) for row in rows))

    # 3. An async sleep is a suspension point by construction.
    await asyncio.sleep(5)
    return response, digest

The second case deserves a caveat: asyncio.to_thread makes the caller cancellable, not the work. The thread runs to completion whatever happens, because Python cannot interrupt a thread. For a long CPU loop the honest fix is to chunk it and check a flag between chunks, or to move it to a process pool where it can be terminated.

A simple test catches the whole category: cancel the task and assert it finishes within a short bound.

import asyncio

import pytest


async def test_cancellation_is_prompt(worker):
    task = asyncio.create_task(worker.run())
    await worker.started.wait()

    task.cancel()
    # If the body has an uncancellable stretch, this wait_for times out and
    # the test names the problem precisely.
    with pytest.raises(asyncio.CancelledError):
        await asyncio.wait_for(task, timeout=0.5)

Half a second is a deliberate choice: long enough that a healthy unwinding always completes, short enough that a blocking call of any realistic length fails it. Adding this one test per long-running coroutine costs very little and converts "shutdown sometimes takes thirty seconds" from an operational mystery into a named test failure.

The same bound is worth applying to the application's own shutdown routine, which is usually a task group or a gather over every background worker. A shutdown test that cancels the group and asserts it completes within a second exercises every worker's cancellation path at once, and it fails the moment any of them grows a blocking call — which is far earlier and far cheaper than discovering it during a deploy.

Shielding the part that must complete

Some cleanup genuinely must finish: flushing a buffer, releasing a distributed lock, writing an audit record. Shielding it protects it from the in-flight cancellation, and bounding the shield protects the shield from becoming a new hang.

import asyncio


async def stream_and_flush(buffer, connection):
    try:
        await stream_forever(buffer, connection)
    finally:
        # Shielded so the flush completes even though we are being cancelled,
        # and bounded so a stuck flush cannot outlive the cancellation forever.
        async with asyncio.timeout(2.0):
            await asyncio.shield(buffer.flush(connection))

import asyncio

import pytest


async def test_flush_completes_despite_cancellation(buffer, connection):
    task = asyncio.create_task(stream_and_flush(buffer, connection))
    await buffer.first_write.wait()

    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task

    assert buffer.flushed_bytes > 0, "the shielded flush did not run"
    assert connection.closed, "the connection was not closed after the flush"


async def test_stuck_flush_does_not_outlive_the_cancellation(buffer, stuck_connection):
    task = asyncio.create_task(stream_and_flush(buffer, stuck_connection))
    await buffer.first_write.wait()

    task.cancel()
    with pytest.raises((asyncio.CancelledError, TimeoutError)):
        await asyncio.wait_for(task, timeout=5)   # must not hang here

The second test is the one that justifies the bound. Without the asyncio.timeout around the shield, a flush against an unresponsive connection would keep the task alive indefinitely, and the only symptom would be a shutdown that never completes — a failure that is extremely hard to attribute in production and trivial to catch here.

Cancellation in groups and gathers

The unit of cancellation differs between the two composition primitives, and tests need to match.

With asyncio.gather(..., return_exceptions=False), cancelling the gathering task propagates cancellation to the children, but a child raising does not cancel its siblings — they keep running unsupervised. A test that asserts "when one fails the others stop" will fail against gather, correctly, because gather makes no such promise.

With a task group, the first failure cancels every sibling and the block raises an ExceptionGroup. The assertion shape changes accordingly, and the details are in testing code that uses task groups.

Frequently Asked Questions

Why does my coroutine ignore task.cancel()? Either it never reaches an await, because it is blocked in synchronous code, or it catches the CancelledError and does not re-raise. Cancellation is delivered as an exception at the next suspension point, so a CPU loop or a blocking socket read is immune to it until it yields.

Should cleanup be shielded? Only the part that must complete, and only with its own deadline. asyncio.shield around a short release or rollback is legitimate; shielding a whole cleanup coroutine means a timeout can no longer stop it, which reintroduces the hang the timeout existed to prevent.

Is catching CancelledError ever correct? Yes, to run cleanup — and then it must be re-raised. Swallowing it converts a cancellation into a silent delay and breaks every caller that expects the task to stop. From Python 3.8 it inherits from BaseException specifically so a bare except Exception does not catch it by accident.

Related

• Timeouts, Cancellation & Deadlines — the layered deadlines that trigger cancellation in production.
• Testing Async Generators and Context Managers — the same cleanup questions for early exit rather than cancellation.
• Testing Code That Uses Task Groups — asserting sibling cancellation.
• Diagnosing "Task was destroyed" Warnings — what an uncancelled, unreferenced task looks like at shutdown.

← Back to Timeouts, Cancellation & Deadlines