Coverage for zenmqtt/mqtt/protocol/commands.py: 95%

1import asyncio 

2from dataclasses import dataclass 

3from logging import getLogger 

4from typing import Any, Optional, Sequence, TypeVar, Union 

5 

6from zenmqtt.exceptions import NotConnectedError, ServerLimitError 

7from zenmqtt.mqtt.protocol.context import ProtocolContext 

8from zenmqtt.mqtt.subscribe import ( 

9 SubscribeResult, 

10 SubscriptionProperties, 

11 SubscriptionRequest, 

12 UnsubscribeProperties, 

13 UnsubscribeResult, 

14 pack_subscribe_packet, 

15 pack_unsubscribe_packet, 

16) 

17from zenmqtt.session import MQTTSession, PacketIdentifier 

18 

19logger = getLogger(__name__) 

20 

21CommandResult = Union[SubscribeResult, UnsubscribeResult] 

22 

23_Result = TypeVar("_Result", SubscribeResult, UnsubscribeResult) 

24 

25 

26@dataclass(slots=True) 

27class _PendingCommand: 

28 # the future of SUBACK or UNSUBACK 

29 future: asyncio.Future[Any] 

30 sent_at: float 

31 

32 

33class Commands: 

34 """SUBSCRIBE and UNSUBSCRIBE, which wait for SUBACK and UNSUBACK.""" 

35 

36 def __init__(self, context: ProtocolContext, session: MQTTSession) -> None: 

37 self._context = context 

38 self._session = session 

39 

40 self._pending: dict[PacketIdentifier, _PendingCommand] = {} 

41 

42 async def subscribe( 

43 self, 

44 topics: Sequence[SubscriptionRequest], 

45 properties: Optional[SubscriptionProperties] = None, 

46 ) -> SubscribeResult: 

47 self._context.ensure_connected() 

48 

49 properties = properties or {} 

50 

51 self._context.server_limits.check_subscribe(topics, properties) 

52 

53 packet_identifier = await self._session.acquire_packet_identifier() 

54 future: asyncio.Future[SubscribeResult] = self._create_future() 

55 

56 return await self._send( 

57 packet_identifier, 

58 pack_subscribe_packet(packet_identifier, topics, properties), 

59 future, 

60 ) 

61 

62 async def unsubscribe( 

63 self, topics: Sequence[str], properties: Optional[UnsubscribeProperties] = None 

64 ) -> UnsubscribeResult: 

65 self._context.ensure_connected() 

66 

67 properties = properties or {} 

68 

69 packet_identifier = await self._session.acquire_packet_identifier() 

70 future: asyncio.Future[UnsubscribeResult] = self._create_future() 

71 

72 return await self._send( 

73 packet_identifier, 

74 pack_unsubscribe_packet(packet_identifier, topics, properties), 

75 future, 

76 ) 

77 

78 async def handle_suback(self, suback_packet: SubscribeResult) -> None: 

79 logger.debug("mqtt_protocol.handle_suback_packet packet:%s", suback_packet) 

80 

81 await self._complete(suback_packet.packet_identifier, suback_packet) 

82 

83 async def handle_unsuback(self, unsuback_packet: UnsubscribeResult) -> None: 

84 logger.debug("mqtt_protocol.handle_unsuback_packet packet:%s", unsuback_packet) 

85 

86 await self._complete(unsuback_packet.packet_identifier, unsuback_packet) 

87 

88 def on_connection_lost(self, exc: Exception) -> list[PacketIdentifier]: 

89 """ 

90 Fails waiting commands, they aren't re-sent. Returns their packet 

91 identifiers, which must be released. 

92 """ 

93 pending, self._pending = self._pending, {} 

94 

95 for command in pending.values(): 

96 if not command.future.done(): 96 ↛ 95line 96 didn't jump to line 95 because the condition on line 96 was always true

97 command.future.set_exception(exc) 

98 

99 return list(pending) 

100 

101 @staticmethod 

102 def _create_future() -> asyncio.Future[Any]: 

103 return asyncio.get_running_loop().create_future() 

104 

105 async def _send( 

106 self, 

107 packet_identifier: PacketIdentifier, 

108 packet: bytes, 

109 future: asyncio.Future[_Result], 

110 ) -> _Result: 

111 try: 

112 self._context.server_limits.check_packet_size(packet) 

113 connection = self._context.ensure_connected() 

114 except (NotConnectedError, ServerLimitError): 

115 await self._session.release_packet_identifier(packet_identifier) 

116 raise 

117 

118 self._pending[packet_identifier] = _PendingCommand(future, self._context.now()) 

119 

120 await self._context.write(connection, packet) 

121 

122 return await future 

123 

124 async def _complete( 

125 self, packet_identifier: PacketIdentifier, result: CommandResult 

126 ) -> None: 

127 await self._session.release_packet_identifier(packet_identifier) 

128 

129 if (command := self._pending.pop(packet_identifier, None)) is None: 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true

130 return 

131 

132 duration = self._context.now() - command.sent_at 

133 metrics = self._context.metrics 

134 

135 if isinstance(result, SubscribeResult): 

136 metrics.on_subscribe_completed(result.reason_codes, duration) 

137 else: 

138 metrics.on_unsubscribe_completed(result.reason_codes, duration) 

139 

140 if not command.future.done(): 140 ↛ exitline 140 didn't return from function '_complete' because the condition on line 140 was always true

141 command.future.set_result(result)