Coverage for zenmqtt/subscriptions.py: 81%

1from logging import getLogger 

2from typing import Awaitable, Callable, Final, Sequence 

3 

4from zenmqtt.mqtt.reason_codes import FAILURE_REASON_CODE 

5from zenmqtt.mqtt.subscribe import ( 

6 Subscription, 

7 SubscriptionProperties, 

8 SubscriptionRequest, 

9 to_subscription, 

10) 

11 

12logger = getLogger(__name__) 

13 

14# UNSUBACK reason code: the topic wasn't subscribed, it's not an error 

15_NO_SUBSCRIPTION_EXISTED: Final[int] = 0x11 

16 

17Subscribe = Callable[ 

18 [Sequence[Subscription], SubscriptionProperties], Awaitable[object] 

19] 

20 

21 

22class SubscriptionRegistry: 

23 """ 

24 Subscriptions accepted by the server; they are restored after reconnect if 

25 the server lost the session. 

26 """ 

27 

28 def __init__(self) -> None: 

29 self._subscriptions: dict[str, tuple[Subscription, SubscriptionProperties]] = {} 

30 

31 def add( 

32 self, 

33 requests: Sequence[SubscriptionRequest], 

34 properties: SubscriptionProperties, 

35 reason_codes: Sequence[int], 

36 ) -> None: 

37 """Remembers the subscriptions which SUBACK accepted.""" 

38 for request, reason_code in zip(requests, reason_codes): 

39 if reason_code < FAILURE_REASON_CODE: 

40 subscription = to_subscription(request) 

41 self._subscriptions[subscription.topic] = (subscription, properties) 

42 

43 def remove(self, topics: Sequence[str], reason_codes: Sequence[int]) -> None: 

44 """Forgets the subscriptions which UNSUBACK removed.""" 

45 for topic, reason_code in zip(topics, reason_codes): 

46 if ( 46 ↛ 45line 46 didn't jump to line 45 because the condition on line 46 was always true

47 reason_code < FAILURE_REASON_CODE 

48 or reason_code == _NO_SUBSCRIPTION_EXISTED 

49 ): 

50 self._subscriptions.pop(topic, None) 

51 

52 async def restore(self, subscribe: Subscribe) -> None: 

53 """ 

54 Subscribes again, one SUBSCRIBE per group of subscriptions with the same 

55 properties. Errors are logged, the other groups are restored anyway. 

56 """ 

57 groups: list[tuple[SubscriptionProperties, list[Subscription]]] = [] 

58 

59 for subscription, properties in self._subscriptions.values(): 

60 for group_properties, subscriptions in groups: 60 ↛ 61line 60 didn't jump to line 61 because the loop on line 60 never started

61 if group_properties == properties: 

62 subscriptions.append(subscription) 

63 break 

64 else: 

65 groups.append((properties, [subscription])) 

66 

67 for properties, subscriptions in groups: 

68 try: 

69 await subscribe(subscriptions, properties) 

70 except Exception as exc: 

71 logger.error( 

72 "mqtt_client.restore_subscriptions.failed: %s", 

73 [subscription.topic for subscription in subscriptions], 

74 exc_info=exc, 

75 )