From e605d4aed28b99b28662372b8ab3f79cf57b1d57 Mon Sep 17 00:00:00 2001 From: Shreyas G Gowda Date: Sat, 28 Feb 2026 08:51:33 +0530 Subject: [PATCH 1/2] test: add tests for connection error handling --- tests/test_connection_errors.py | 83 +++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 tests/test_connection_errors.py diff --git a/tests/test_connection_errors.py b/tests/test_connection_errors.py new file mode 100644 index 0000000..4713c30 --- /dev/null +++ b/tests/test_connection_errors.py @@ -0,0 +1,83 @@ +"""Tests for connection error handling.""" + +import pytest +from unittest.mock import AsyncMock, patch + +from streamline_sdk import StreamlineClient +from streamline_sdk.exceptions import ConnectionError +from streamline_sdk.producer import Producer +from streamline_sdk.admin import Admin + + +class TestConnectionErrors: + """Tests for connection failure scenarios.""" + + @pytest.mark.asyncio + async def test_server_unreachable(self): + """Test that ConnectionError is raised when server is unreachable.""" + with patch.object(Producer, "start", new_callable=AsyncMock) as mock_start: + mock_start.side_effect = OSError("Connection refused") + + client = StreamlineClient(bootstrap_servers="localhost:9092") + with pytest.raises(ConnectionError, match="Failed to start client"): + await client.start() + + @pytest.mark.asyncio + async def test_timeout_error(self): + """Test that ConnectionError is raised on connection timeout.""" + with patch.object(Producer, "start", new_callable=AsyncMock) as mock_start: + mock_start.side_effect = TimeoutError("Connection timed out") + + client = StreamlineClient(bootstrap_servers="localhost:9092") + with pytest.raises(ConnectionError, match="Failed to start client"): + await client.start() + + @pytest.mark.asyncio + async def test_dns_resolution_failure(self): + """Test that ConnectionError is raised when DNS resolution fails.""" + with patch.object(Producer, "start", new_callable=AsyncMock) as mock_start: + mock_start.side_effect = OSError("Name or service not known") + + client = StreamlineClient(bootstrap_servers="unknownhost:9092") + with pytest.raises(ConnectionError, match="Failed to start client"): + await client.start() + + @pytest.mark.asyncio + async def test_client_closed_after_failed_connection(self): + """Test that client is properly cleaned up after connection failure.""" + with patch.object(Producer, "start", new_callable=AsyncMock) as mock_start: + mock_start.side_effect = OSError("Connection refused") + + client = StreamlineClient(bootstrap_servers="localhost:9092") + with pytest.raises(ConnectionError): + await client.start() + + assert client.is_connected is False + + @pytest.mark.asyncio + async def test_reconnection_after_failure(self): + """Test that client can attempt reconnection after failure.""" + with patch.object(Producer, "start", new_callable=AsyncMock) as mock_start: + mock_start.side_effect = OSError("Connection refused") + + client = StreamlineClient(bootstrap_servers="localhost:9092") + with pytest.raises(ConnectionError): + await client.start() + + # Now try again with successful connection + with patch.object(Producer, "start", new_callable=AsyncMock): + with patch.object(Admin, "start", new_callable=AsyncMock): + client2 = StreamlineClient(bootstrap_servers="localhost:9092") + await client2.start() + assert client2.is_connected is True + await client2.close() + + @pytest.mark.asyncio + async def test_context_manager_handles_connection_error(self): + """Test that async context manager handles connection errors gracefully.""" + with patch.object(Producer, "start", new_callable=AsyncMock) as mock_start: + mock_start.side_effect = OSError("Connection refused") + + with pytest.raises(ConnectionError): + async with StreamlineClient(bootstrap_servers="localhost:9092"): + pass \ No newline at end of file From f4228d153a5eb088cc5c17332e858c6b2c8cc1cf Mon Sep 17 00:00:00 2001 From: Shreyas G Gowda Date: Sat, 28 Feb 2026 16:48:41 +0530 Subject: [PATCH 2/2] fix: add type annotations to all public API methods --- streamline_sdk/admin.py | 2 +- streamline_sdk/client.py | 10 +++++----- streamline_sdk/consumer.py | 11 ++++++----- streamline_sdk/exceptions.py | 25 ++++++++++--------------- streamline_sdk/producer.py | 2 +- streamline_sdk/retry.py | 4 ++-- streamline_sdk/serializers.py | 6 +++--- streamline_sdk/telemetry.py | 14 ++++---------- streamline_sdk/types.py | 6 +++--- 9 files changed, 35 insertions(+), 45 deletions(-) diff --git a/streamline_sdk/admin.py b/streamline_sdk/admin.py index e49049b..94a8b2a 100644 --- a/streamline_sdk/admin.py +++ b/streamline_sdk/admin.py @@ -338,6 +338,6 @@ async def __aenter__(self) -> "Admin": await self.start() return self - async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + async def __aexit__(self, exc_type: Optional[type], exc_val: Optional[BaseException], exc_tb: Optional[Any]) -> None: """Exit async context manager.""" await self.close() diff --git a/streamline_sdk/client.py b/streamline_sdk/client.py index 5786e09..63213e1 100644 --- a/streamline_sdk/client.py +++ b/streamline_sdk/client.py @@ -5,7 +5,7 @@ import asyncio import os from dataclasses import dataclass, field -from typing import Optional +from typing import Optional,Any from .producer import Producer from .consumer import Consumer @@ -113,8 +113,8 @@ def __init__( client_config: Optional[ClientConfig] = None, producer_config: Optional[ProducerConfig] = None, consumer_config: Optional[ConsumerConfig] = None, - **kwargs, - ): + **kwargs: Any, + ) -> None: """Initialize the client. Args: @@ -173,7 +173,7 @@ def admin(self) -> Admin: def consumer( self, group_id: Optional[str] = None, - **kwargs, + **kwargs:Any, ) -> Consumer: """Create a new consumer. @@ -255,7 +255,7 @@ async def __aenter__(self) -> "StreamlineClient": await self.start() return self - async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + async def __aexit__(self, exc_type: Optional[type], exc_val: Optional[BaseException], exc_tb: Optional[Any]) -> None: """Exit the async context manager.""" await self.close() diff --git a/streamline_sdk/consumer.py b/streamline_sdk/consumer.py index 7aa80a8..213a4e8 100644 --- a/streamline_sdk/consumer.py +++ b/streamline_sdk/consumer.py @@ -200,7 +200,7 @@ async def position(self, partition: TopicPartition) -> int: if self._consumer is None: raise ConsumerError("Consumer not started") - return await self._consumer.position(partition) + return int(await self._consumer.position(partition)) async def committed(self, partition: TopicPartition) -> Optional[int]: """Get committed offset for a partition. @@ -214,7 +214,8 @@ async def committed(self, partition: TopicPartition) -> Optional[int]: if self._consumer is None: raise ConsumerError("Consumer not started") - return await self._consumer.committed(partition) + result = await self._consumer.committed(partition) + return int(result) if result is not None else None def assignment(self) -> Set[TopicPartition]: """Get assigned partitions. @@ -225,7 +226,7 @@ def assignment(self) -> Set[TopicPartition]: if self._consumer is None: return set() - return self._consumer.assignment() + return set(self._consumer.assignment()) def subscription(self) -> Set[str]: """Get subscribed topics. @@ -312,13 +313,13 @@ def is_started(self) -> bool: @property def group_id(self) -> Optional[str]: """Get the consumer group ID.""" - return self._consumer_config.group_id + return str(self._consumer_config.group_id) if self._consumer_config.group_id is not None else None async def __aenter__(self) -> "Consumer": """Enter async context manager.""" await self.start() return self - async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + async def __aexit__(self, exc_type: Optional[type], exc_val: Optional[BaseException], exc_tb: Optional[Any]) -> None: """Exit async context manager.""" await self.close() diff --git a/streamline_sdk/exceptions.py b/streamline_sdk/exceptions.py index 420ceb9..24abe94 100644 --- a/streamline_sdk/exceptions.py +++ b/streamline_sdk/exceptions.py @@ -1,14 +1,15 @@ """Exceptions for the Streamline SDK.""" +from typing import Any, Optional class StreamlineError(Exception): """Base exception for Streamline SDK errors.""" - def __init__(self, *args, hint: str | None = None, **kwargs): + def __init__(self, *args: Any, hint: Optional[str] = None, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.hint = hint - def __str__(self): + def __str__(self) -> str: s = super().__str__() if self.hint: s += f" (hint: {self.hint})" @@ -18,7 +19,7 @@ def __str__(self): class ConnectionError(StreamlineError): """Error connecting to Streamline.""" - def __init__(self, *args, hint: str | None = None, **kwargs): + def __init__(self, *args: Any, hint: Optional[str] = None, **kwargs: Any) -> None: super().__init__( *args, hint=hint or "Check that Streamline server is running and accessible", @@ -28,24 +29,21 @@ def __init__(self, *args, hint: str | None = None, **kwargs): class ProducerError(StreamlineError): """Error in producer operations.""" - pass class ConsumerError(StreamlineError): """Error in consumer operations.""" - pass class TopicError(StreamlineError): """Error in topic operations.""" - def __init__(self, *args, hint: str | None = None, **kwargs): + def __init__(self, *args: Any, hint: Optional[str] = None, **kwargs: Any) -> None: super().__init__( *args, - hint=hint - or "Use admin client to create the topic first, or enable auto-creation", + hint=hint or "Use admin client to create the topic first, or enable auto-creation", **kwargs, ) @@ -53,7 +51,7 @@ def __init__(self, *args, hint: str | None = None, **kwargs): class AuthenticationError(StreamlineError): """Error in authentication.""" - def __init__(self, *args, hint: str | None = None, **kwargs): + def __init__(self, *args: Any, hint: Optional[str] = None, **kwargs: Any) -> None: super().__init__( *args, hint=hint or "Verify your SASL credentials and mechanism", @@ -63,23 +61,20 @@ def __init__(self, *args, hint: str | None = None, **kwargs): class AuthorizationError(StreamlineError): """Error in authorization (ACL denied).""" - pass class SerializationError(StreamlineError): """Error in message serialization/deserialization.""" - pass class TimeoutError(StreamlineError): """Operation timed out.""" - def __init__(self, *args, hint: str | None = None, **kwargs): + def __init__(self, *args: Any, hint: Optional[str] = None, **kwargs: Any) -> None: super().__init__( *args, - hint=hint - or "Consider increasing timeout settings or checking server load", + hint=hint or "Consider increasing timeout settings or checking server load", **kwargs, - ) + ) \ No newline at end of file diff --git a/streamline_sdk/producer.py b/streamline_sdk/producer.py index 1f4623f..4446d24 100644 --- a/streamline_sdk/producer.py +++ b/streamline_sdk/producer.py @@ -232,6 +232,6 @@ async def __aenter__(self) -> "Producer": await self.start() return self - async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + async def __aexit__(self, exc_type: Optional[type], exc_val: Optional[BaseException], exc_tb: Optional[Any]) -> None: """Exit async context manager.""" await self.close() diff --git a/streamline_sdk/retry.py b/streamline_sdk/retry.py index d03e5a4..0bfa189 100644 --- a/streamline_sdk/retry.py +++ b/streamline_sdk/retry.py @@ -110,7 +110,7 @@ async def retry_async( raise last_exception # type: ignore[misc] -def with_retry(config: Optional[RetryConfig] = None) -> Callable: +def with_retry(config: Optional[RetryConfig] = None) -> Callable[..., Any]: """Decorator that adds retry logic to an async function. Args: @@ -127,7 +127,7 @@ async def send_message(topic, value): """ cfg = config or RetryConfig() - def decorator(func: Callable) -> Callable: + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(func) async def wrapper(*args: Any, **kwargs: Any) -> Any: return await retry_async(func, *args, config=cfg, **kwargs) diff --git a/streamline_sdk/serializers.py b/streamline_sdk/serializers.py index 72083af..7571098 100644 --- a/streamline_sdk/serializers.py +++ b/streamline_sdk/serializers.py @@ -66,7 +66,7 @@ async def register_schema(self, subject: str, schema_str: str, schema_type: str data = await resp.json() schema_id = data["id"] self._cache[subject] = schema_id - return schema_id + return int(schema_id) else: text = await resp.text() raise StreamlineError(f"Schema registration failed: {resp.status} {text}") @@ -79,7 +79,7 @@ async def get_schema(self, schema_id: int) -> str: async with session.get(url) as resp: if resp.status == 200: data = await resp.json() - return data["schema"] + return str(data["schema"]) else: raise StreamlineError(f"Schema not found: {schema_id}") @@ -151,7 +151,7 @@ def __init__( self.auto_register = auto_register self.validate = validate self._schema_id: Optional[int] = None - self._schema: Optional[dict] = json.loads(schema_str) if schema_str else None + self._schema: Optional[Dict[str, Any]] = json.loads(schema_str) if schema_str else None async def serialize(self, topic: str, value: Dict[str, Any]) -> bytes: """Serialize a value to JSON bytes with optional validation.""" diff --git a/streamline_sdk/telemetry.py b/streamline_sdk/telemetry.py index eeea83e..930d693 100644 --- a/streamline_sdk/telemetry.py +++ b/streamline_sdk/telemetry.py @@ -316,10 +316,7 @@ def trace_consume_sync( # ── Decorators ──────────────────────────────────────────────────── - def traced_produce( - self, - topic: str, - ) -> Callable: + def traced_produce(self, topic: str) -> Callable[..., Any]: """Decorator that traces an async produce function. Args: @@ -332,7 +329,7 @@ async def publish(data): await producer.send("events", value=data) """ - def decorator(func: Callable) -> Callable: + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: if not self._enabled: return func @@ -345,10 +342,7 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any: return decorator - def traced_consume( - self, - topic: str, - ) -> Callable: + def traced_consume(self, topic: str) -> Callable[..., Any]: """Decorator that traces an async consume function. Args: @@ -362,7 +356,7 @@ async def handle(records): process(record) """ - def decorator(func: Callable) -> Callable: + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: if not self._enabled: return func diff --git a/streamline_sdk/types.py b/streamline_sdk/types.py index 8d0a47e..e0e5ced 100644 --- a/streamline_sdk/types.py +++ b/streamline_sdk/types.py @@ -3,7 +3,7 @@ """ from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional,Iterator from datetime import datetime @@ -224,8 +224,8 @@ class QueryResult: execution_time_ms: int """Execution time in milliseconds.""" - def __iter__(self): + def __iter__(self) -> "Iterator[Dict[str, Any]]": return iter(self.rows) - def __len__(self): + def __len__(self) -> int: return self.row_count