test_shield.py

import asyncio

import pytest


class Mutable:
    def __init__(self) -> None:
        self._mutations: int = 0
    
    @property
    def mutations(self) -> int:
        return self._mutations
    
    def mutate(self) -> None:
        self._mutations += 1
    
    async def a_mutate(self, delay: float = 0) -> None:
        await asyncio.sleep(delay)
        self._mutations += 1


async def wait_then_call(f, delay: float = 0.2):
    await asyncio.sleep(delay)
    await f()


async def wait_for_cancellation(f, delay: float = 0.2, reraise = False):
    await asyncio.sleep(delay)
    try:
        await f
    except asyncio.CancelledError:
        if reraise:
            raise
    else:
        raise AssertionError(f"Did not raise {asyncio.CancelledError}")


def test_sync_mutations():
    m = Mutable()
    n = m
    assert n.mutations == 0
    n.mutate()
    assert m.mutations == 1
    m.mutate()
    assert n.mutations == 2


@pytest.mark.asyncio
async def test_async_mutations():
    m = Mutable()
    t = asyncio.create_task(m.a_mutate(delay=0.5))
    await asyncio.sleep(0.2)
    await asyncio.sleep(0.5)
    await t
    assert m.mutations == 1


@pytest.mark.asyncio
async def test_cancellation():
    m = Mutable()
    t = asyncio.create_task(m.a_mutate(delay=0.5))
    with pytest.raises(asyncio.CancelledError):
        await asyncio.sleep(0.2)
        t.cancel()
        await asyncio.sleep(0.5)
        await t
    assert m.mutations == 0


@pytest.mark.asyncio
async def test_deep_cancellation():
    """
    w2 will encounter a CancelledError -- caused by 
    t's cancellation -- and re-raise it to w2, who
    will handle it silently.
    """
    w2 = wait_for_cancellation(asyncio.sleep(0.3), reraise=True)
    w1 = wait_for_cancellation(w2)
    t = asyncio.create_task(w1)
    await asyncio.sleep(0.5)
    t.cancel()
    await t


@pytest.mark.asyncio
async def test_shielding():
    """
    t's cancellation will _not_ cause s to encounter a
    CancelledError, but w will catch a CancelledError at
    s's call-site.
    
    when awaited, s will finish their action, mutating m.
    """
    m = Mutable()
    s = asyncio.shield(m.a_mutate(delay=0.2))
    w = wait_for_cancellation(s)
    t = asyncio.create_task(w)
    await asyncio.sleep(0.1)
    t.cancel()
    assert m.mutations == 0
    await s
    assert m.mutations == 1


@pytest.mark.asyncio
async def test_shielding_propagation():
    """
    t's cancellation will not cause s to encounter a
    CancelledError, and since s is awaiting u, u won't
    be cancelled either. w will handle the cancellation
    from s's callsite.
    """
    m = Mutable()
    u = m.a_mutate(delay=0.2)
    s = asyncio.shield(asyncio.gather(m.a_mutate(delay=0.2), u))
    w = wait_for_cancellation(s)
    t = asyncio.create_task(w)
    await asyncio.sleep(0.1)
    t.cancel()
    assert m.mutations == 0
    await s
    assert m.mutations == 2