diff options
Diffstat (limited to 'tests/test_asyncio/test_pubsub.py')
| -rw-r--r-- | tests/test_asyncio/test_pubsub.py | 22 |
1 files changed, 13 insertions, 9 deletions
diff --git a/tests/test_asyncio/test_pubsub.py b/tests/test_asyncio/test_pubsub.py index c2a9130..0df7847 100644 --- a/tests/test_asyncio/test_pubsub.py +++ b/tests/test_asyncio/test_pubsub.py @@ -5,7 +5,11 @@ import sys from typing import Optional from unittest.mock import patch -import async_timeout +if sys.version_info.major >= 3 and sys.version_info.minor >= 11: + from asyncio import timeout as async_timeout +else: + from async_timeout import timeout as async_timeout + import pytest import pytest_asyncio @@ -21,7 +25,7 @@ def with_timeout(t): def wrapper(corofunc): @functools.wraps(corofunc) async def run(*args, **kwargs): - async with async_timeout.timeout(t): + async with async_timeout(t): return await corofunc(*args, **kwargs) return run @@ -648,7 +652,7 @@ class TestPubSubReconnect: async def loop(): # must make sure the task exits - async with async_timeout.timeout(2): + async with async_timeout(2): nonlocal interrupt await pubsub.subscribe("foo") while True: @@ -677,7 +681,7 @@ class TestPubSubReconnect: task = asyncio.get_running_loop().create_task(loop()) # get the initial connect message - async with async_timeout.timeout(1): + async with async_timeout(1): message = await messages.get() assert message == { "channel": b"foo", @@ -776,7 +780,7 @@ class TestPubSubRun: if n == 1: break await asyncio.sleep(0.1) - async with async_timeout.timeout(0.1): + async with async_timeout(0.1): message = await messages.get() task.cancel() # we expect a cancelled error, not the Runtime error @@ -839,7 +843,7 @@ class TestPubSubAutoReconnect: Test that a socket error will cause reconnect """ try: - async with async_timeout.timeout(self.timeout): + async with async_timeout(self.timeout): await self.mysetup(r, method) # now, disconnect the connection, and wait for it to be re-established async with self.cond: @@ -868,7 +872,7 @@ class TestPubSubAutoReconnect: Test that a manual disconnect() will cause reconnect """ try: - async with async_timeout.timeout(self.timeout): + async with async_timeout(self.timeout): await self.mysetup(r, method) # now, disconnect the connection, and wait for it to be re-established async with self.cond: @@ -923,7 +927,7 @@ class TestPubSubAutoReconnect: async def loop_step_listen(self): # get a single message via listen() try: - async with async_timeout.timeout(0.1): + async with async_timeout(0.1): async for message in self.pubsub.listen(): await self.messages.put(message) return True @@ -947,7 +951,7 @@ class TestBaseException: assert pubsub.connection.is_connected async def get_msg_or_timeout(timeout=0.1): - async with async_timeout.timeout(timeout): + async with async_timeout(timeout): # blocking method to return messages while True: response = await pubsub.parse_response(block=True) |