fix(TestStreamClient): missing utility TimeoutErrorException added to…

kstreams/test_utils/test_clients.py

@@ -18,7 +18,7 @@
 from kstreams.clients import Consumer, Producer
 from kstreams.serializers import NO_DEFAULT, Serializer
 from kstreams.types import ConsumerRecord, EncodedHeaders
-from tests import TimeoutErrorException
+from kstreams.utils import TimeoutErrorException
 
 from .structs import RecordMetadata
 from .topics import TopicManager

kstreams/utils.py

@@ -5,7 +5,7 @@
 import ssl
 import sys
 from tempfile import NamedTemporaryFile
-from typing import Any, Optional, Union
+from typing import Any, Optional, Type, Union
 
 from aiokafka.helpers import create_ssl_context as aiokafka_create_ssl_context
 
@@ -16,6 +16,14 @@
 PY_VERSION = sys.version_info
 
 
+if PY_VERSION < (3, 11):
+    TimeoutErrorException: Type[Union[asyncio.TimeoutError, TimeoutError]] = (
+        asyncio.TimeoutError
+    )
+else:
+    TimeoutErrorException = TimeoutError
+
+
 def encode_headers(headers: types.Headers) -> types.EncodedHeaders:
     return [(header, value.encode()) for header, value in headers.items()]
 

tests/__init__.py

@@ -1,7 +0,0 @@
-import asyncio
-import sys
-
-if sys.version_info < (3, 11):
-    TimeoutErrorException = asyncio.TimeoutError
-else:
-    TimeoutErrorException = TimeoutError

tests/test_client.py

@@ -23,7 +23,7 @@
     TestStreamClient,
     TopicManager,
 )
-from tests import TimeoutErrorException
+from kstreams.utils import TimeoutErrorException
 
 topic = "local--kstreams-consumer"
 tp0 = TopicPartition(topic=topic, partition=0)

tests/test_streams.py

@@ -10,7 +10,7 @@
 from kstreams.streams import stream
 from kstreams.structs import TopicPartitionOffset
 from kstreams.test_utils import TestStreamClient
-from tests import TimeoutErrorException
+from kstreams.utils import TimeoutErrorException
 
 # NOTE: remove the test when `no typing` support is deprecated