Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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
Add Comment
Please, Sign In to add comment