test: add tests for connection error handling

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()

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()
 

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()

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,32 +29,29 @@ 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,
         )
 
 
 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,
-        )
+        )

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()

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)

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."""

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
 

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

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