test: cover streaming safeguards

Author: AKKI0511

config/streaming.yaml

@@ -4,6 +4,12 @@ streaming:
       websocket_url: "wss://stream.data.alpaca.markets/v2/iex"
       auth_method: "api_key"
       subscriptions: ["trades", "quotes"]
+      rate_limit:
+        default_rate: 100
+        burst_allowance: 50
+      circuit_breaker:
+        failure_threshold: 5
+        timeout: 30
   buffer_size: 10000
   reconnect_attempts: 5
   health_check_interval: 30

poetry.lock

@@ -26,7 +26,11 @@ dependencies = [
     "pydantic (>=2.11.7,<3.0.0)",
     "alpha-vantage (>=3.0.0,<4.0.0)",
     "litellm (>=1.75.0,<2.0.0)",
-    "websockets (>=15.0.1,<16.0.0)"
+    "websockets (>=15.0.1,<16.0.0)",
+    "structlog (>=25.4.0,<26.0.0)",
+    "pybreaker (>=1.4.0,<2.0.0)",
+    "asyncio-throttle (>=1.0.2,<2.0.0)",
+    "prometheus-client (>=0.22.1,<0.23.0)"
 ]
 
 [tool.poetry]

pyproject.toml

@@ -1,5 +1,13 @@
 """Streaming infrastructure package."""
 
 from .gateway import StreamingGateway
+from .auth_manager import AuthManager
+from .rate_limiter import AdaptiveRateLimiter
+from .connection_pool import ConnectionPool
 
-__all__ = ["StreamingGateway"]
+__all__ = [
+    "StreamingGateway",
+    "AuthManager",
+    "AdaptiveRateLimiter",
+    "ConnectionPool",
+]

quanttradeai/streaming/__init__.py

@@ -9,6 +9,9 @@
 
 import websockets
 
+from ..auth_manager import AuthManager
+from ..rate_limiter import AdaptiveRateLimiter
+
 
 @dataclass
 class DataProviderAdapter(ABC):
@@ -19,18 +22,28 @@ class DataProviderAdapter(ABC):
     connection: Optional[websockets.WebSocketClientProtocol] = field(
         default=None, init=False
     )
+    circuit_breaker: Any = field(default=None, init=False)
+    rate_limiter: Optional[AdaptiveRateLimiter] = field(default=None, init=False)
+    auth_manager: Optional[AuthManager] = field(default=None, init=False)
 
     async def connect(self) -> None:
         """Establish a WebSocket connection to the provider."""
 
-        self.connection = await websockets.connect(self.websocket_url)
+        headers = None
+        if self.auth_manager:
+            headers = await self.auth_manager.get_auth_headers()
+        self.connection = await websockets.connect(
+            self.websocket_url, extra_headers=headers
+        )
 
     async def subscribe(self, channel: str, symbols: List[str]) -> None:
         """Send a subscription message for ``symbols`` on ``channel``."""
 
         if self.connection is None:
             await self.connect()
         message = self._build_subscribe_message(channel, symbols)
+        if self.rate_limiter:
+            await self.rate_limiter.acquire()
         await self.connection.send(json.dumps(message))
 
     async def listen(self) -> AsyncIterator[Dict[str, Any]]:

quanttradeai/streaming/adapters/base_adapter.py

@@ -0,0 +1,37 @@
+"""Authentication and authorization utilities for data providers."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Dict, Optional
+
+
+@dataclass
+class AuthManager:
+    """Simple authentication manager with token refresh support."""
+
+    provider: str
+    _token: Optional[str] = None
+    _expires_at: datetime = datetime.fromtimestamp(0)
+
+    def _load_credentials(self) -> Dict[str, str]:
+        """Load credentials from environment variables."""
+        key = os.getenv(f"{self.provider.upper()}_API_KEY", "")
+        secret = os.getenv(f"{self.provider.upper()}_API_SECRET", "")
+        return {"key": key, "secret": secret}
+
+    def _token_needs_refresh(self) -> bool:
+        return datetime.utcnow() >= self._expires_at - timedelta(minutes=5)
+
+    async def _refresh_token(self) -> None:
+        creds = self._load_credentials()
+        # In production, call provider-specific auth endpoint.
+        self._token = creds.get("key")
+        self._expires_at = datetime.utcnow() + timedelta(hours=1)
+
+    async def get_auth_headers(self) -> Dict[str, str]:
+        if self._token is None or self._token_needs_refresh():
+            await self._refresh_token()
+        return {"Authorization": f"Bearer {self._token}"}

quanttradeai/streaming/auth_manager.py

@@ -0,0 +1,44 @@
+"""Connection pooling and health management."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Any, Set, Callable, Awaitable
+
+
+@dataclass
+class ConnectionPool:
+    """Maintain a pool of reusable WebSocket connections."""
+
+    max_connections: int = 10
+    _pool: asyncio.Queue = field(init=False)
+    _active: Set[Any] = field(init=False, default_factory=set)
+
+    def __post_init__(self) -> None:
+        self._pool = asyncio.Queue(maxsize=self.max_connections)
+
+    async def acquire_connection(self, factory: Callable[[], Awaitable[Any]]) -> Any:
+        if not self._pool.empty():
+            conn = await self._pool.get()
+        elif len(self._active) < self.max_connections:
+            conn = await factory()
+        else:
+            conn = await self._pool.get()
+        self._active.add(conn)
+        return conn
+
+    async def release(self, conn: Any) -> None:
+        if conn in self._active:
+            self._active.remove(conn)
+            await self._pool.put(conn)
+
+    async def _ping_all_connections(self) -> None:
+        for conn in list(self._active):
+            if hasattr(conn, "ping"):
+                await conn.ping()
+
+    async def health_check_loop(self, interval: int = 30) -> None:
+        while True:
+            await self._ping_all_connections()
+            await asyncio.sleep(interval)

quanttradeai/streaming/connection_pool.py

@@ -14,6 +14,7 @@
 from .adapters.ib_adapter import IBAdapter
 from .processors import MessageProcessor
 from .monitoring.metrics import Metrics
+from .logging import logger
 
 AdapterMap = {
     "alpaca": AlpacaAdapter,
@@ -51,7 +52,12 @@ def __post_init__(self) -> None:
             if adapter_cls is None:
                 raise ValueError(f"Unknown provider: {name}")
             adapter = adapter_cls(websocket_url=url)
-            self.websocket_manager.add_adapter(adapter)
+            self.websocket_manager.add_adapter(
+                adapter,
+                circuit_breaker_cfg=provider_cfg.get("circuit_breaker", {}),
+                rate_limit_cfg=provider_cfg.get("rate_limit"),
+                auth_method=provider_cfg.get("auth_method", "api_key"),
+            )
 
     # Subscription API -------------------------------------------------
     def subscribe_to_trades(self, symbols: List[str], callback: Callback) -> None:
@@ -68,10 +74,14 @@ def subscribe_to_quotes(self, symbols: List[str], callback: Callback) -> None:
 
     # Runtime ----------------------------------------------------------
     async def _dispatch(self, provider: str, message: Dict) -> None:
-        self.metrics.increment()
+        symbol = message.get("symbol", "")
+        self.metrics.record_message(provider, symbol)
         processed = self.message_processor.process(message)
         await self.buffer.put(processed)
         msg_type = message.get("type", "trades")
+        logger.debug(
+            "dispatch_message", provider=provider, type=msg_type, symbol=symbol
+        )
         for cb in self._callbacks.get(msg_type, []):
             res = cb(processed)
             if asyncio.iscoroutine(res):

quanttradeai/streaming/gateway.py

@@ -0,0 +1,27 @@
+"""Structured logging utilities for streaming infrastructure."""
+
+from __future__ import annotations
+
+import logging
+import structlog
+
+
+def configure_logging(level: str = "INFO") -> None:
+    """Configure structlog for JSON formatted logs.
+
+    Args:
+        level: Minimum log level.
+    """
+    structlog.configure(
+        processors=[
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.processors.add_log_level,
+            structlog.contextvars.merge_contextvars,
+            structlog.processors.JSONRenderer(),
+        ],
+        wrapper_class=structlog.make_filtering_bound_logger(getattr(logging, level)),
+    )
+
+
+configure_logging()
+logger = structlog.get_logger()

quanttradeai/streaming/logging.py

@@ -1,17 +1,30 @@
-"""Streaming metrics collection."""
+"""Prometheus metrics for streaming infrastructure."""
 
 from __future__ import annotations
 
 from dataclasses import dataclass
 
+from prometheus_client import Counter, Gauge, Histogram
+
+
+messages_processed = Counter(
+    "websocket_messages_total", "Total processed messages", ["provider", "symbol"]
+)
+connection_latency = Histogram(
+    "websocket_connection_latency_seconds", "Connection establishment time"
+)
+active_connections = Gauge("websocket_active_connections", "Current active connections")
+
 
 @dataclass
 class Metrics:
-    """Simple in-memory metrics tracker."""
+    """Convenience wrapper around Prometheus metrics."""
 
-    messages_received: int = 0
+    def record_message(self, provider: str, symbol: str) -> None:
+        messages_processed.labels(provider=provider, symbol=symbol).inc()
 
-    def increment(self) -> None:
-        """Increment the received message counter."""
+    def record_connection_latency(self, seconds: float) -> None:
+        connection_latency.observe(seconds)
 
-        self.messages_received += 1
+    def set_active_connections(self, count: int) -> None:
+        active_connections.set(count)