Coverage for src/lexigram/graphql/subscriptions/auth.py: 50%

1"""Subscription authentication hooks. 

2 

3Provides connection-level authentication for WebSocket subscriptions. 

4""" 

5 

6from __future__ import annotations 

7 

8from collections.abc import Awaitable, Callable 

9from typing import Any 

10 

11from lexigram.graphql.exceptions import AuthenticationError 

12 

13# Type for connection auth hook 

14ConnectionAuthHook = Callable[[dict[str, Any]], Awaitable[bool]] 

15 

16 

17class SubscriptionAuth: 

18 """Handles authentication for WebSocket subscriptions. 

19 

20 Provides hooks for validating connection tokens during connection_init. 

21 """ 

22 

23 def __init__( 

24 self, 

25 on_auth: ConnectionAuthHook | None = None, 

26 ): 

27 """Initialize the subscription auth handler. 

28 

29 Args: 

30 on_auth: Optional async callback to validate auth payload. 

31 Should return True if authentication succeeds. 

32 """ 

33 self._on_auth = on_auth 

34 

35 async def authenticate(self, payload: dict[str, Any]) -> bool: 

36 """Authenticate a connection_init payload. 

37 

38 Args: 

39 payload: The payload from the connection_init message. 

40 

41 Returns: 

42 True if authentication succeeds. 

43 

44 Raises: 

45 AuthenticationError: If authentication fails. 

46 """ 

47 if self._on_auth is None: 

48 # No auth hook - allow all connections 

49 return True 

50 

51 try: 

52 result = await self._on_auth(payload) 

53 if not result: 

54 raise AuthenticationError("Authentication failed") 

55 return result 

56 except AuthenticationError: 

57 raise 

58 except (TypeError, ValueError, RuntimeError) as e: 

59 raise AuthenticationError(f"Authentication error: {e}") from e 

60 

61 def set_auth_hook(self, hook: ConnectionAuthHook) -> None: 

62 """Set the authentication hook. 

63 

64 Args: 

65 hook: Async callback to validate auth payload. 

66 """ 

67 self._on_auth = hook 

68 

69 

70# Default auth instance 

71default_auth = SubscriptionAuth() 

72 

73 

74__all__ = ["ConnectionAuthHook", "SubscriptionAuth", "default_auth"]