Coverage for src/lexigram/graphql/subscriptions/transport/_connection.py: 94%

1"""GraphQL WebSocket Transport for subscriptions. 

2 

3Provides WebSocket transport implementation using graphql-transport-ws protocol. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9from dataclasses import dataclass, field 

10from typing import TYPE_CHECKING, Any, cast 

11 

12from lexigram import serialization as json 

13from lexigram.graphql.subscriptions.protocol import GQLWSMessageType 

14from lexigram.graphql.types import SubscriptionInfo 

15from lexigram.logging import get_logger 

16 

17if TYPE_CHECKING: 

18 from collections.abc import AsyncIterator, Awaitable, Callable 

19 

20 from lexigram.contracts.graphql.protocols import SubscriptionAuthHandlerProtocol 

21 from lexigram.contracts.web import WebSocketProtocol 

22 

23logger = get_logger(__name__) 

24 

25 

26@dataclass 

27class SubscriptionConnection: 

28 """Manages active subscriptions for a WebSocket connection.""" 

29 

30 subscriptions: dict[str, SubscriptionInfo] = field(default_factory=dict) 

31 

32 def add(self, subscription_id: str, info: SubscriptionInfo) -> None: 

33 """Add a subscription.""" 

34 self.subscriptions[subscription_id] = info 

35 

36 def remove(self, subscription_id: str) -> None: 

37 """Remove a subscription.""" 

38 self.subscriptions.pop(subscription_id, None) 

39 

40 def get(self, subscription_id: str) -> SubscriptionInfo | None: 

41 """Get a subscription by ID.""" 

42 return self.subscriptions.get(subscription_id)