Coverage for faststream / mqtt / subscriber / usecase.py: 89%

1import warnings (empty)

2from abc import abstractmethod (empty)

3from collections.abc import AsyncIterator, Sequence (empty)

4from contextlib import suppress (empty)

5from typing import TYPE_CHECKING, Any (empty)

6 

7import anyio (empty)

8import zmqtt (empty)

9from typing_extensions import override (empty)

10 

11from faststream._internal.endpoint.subscriber import SubscriberUsecase (empty)

12from faststream._internal.endpoint.subscriber.mixins import ConcurrentMixin, TasksMixin (empty)

13from faststream._internal.endpoint.utils import process_msg (empty)

14from faststream.middlewares import AckPolicy (empty)

15from faststream.mqtt.parser import MQTTBaseParser, MQTTParserV5, MQTTParserV311 (empty)

16from faststream.mqtt.publisher.fake import MQTTFakePublisher (empty)

17 

18if TYPE_CHECKING: (empty)

19 from faststream._internal.endpoint.publisher import PublisherProto 

20 from faststream._internal.endpoint.subscriber import SubscriberSpecification 

21 from faststream._internal.endpoint.subscriber.call_item import CallsCollection 

22 from faststream.message import StreamMessage 

23 from faststream.mqtt.broker.config import MQTTBrokerConfig 

24 from faststream.mqtt.message import MQTTMessage 

25 from faststream.mqtt.subscriber.config import MQTTSubscriberConfig 

26 

27 

28class MQTTBaseSubscriber(TasksMixin, SubscriberUsecase[zmqtt.Message]): (empty)

29 """Base class for all MQTT subscribers.""" 

30 

31 _outer_config: "MQTTBrokerConfig" (empty)

32 

33 def __init__( (empty)

34 self, 

35 config: "MQTTSubscriberConfig", 

36 specification: "SubscriberSpecification[Any, Any]", 

37 calls: "CallsCollection[zmqtt.Message]", 

38 ) -> None: 

39 # version may not be available yet when subscriber is created on a router 

40 # before include_router is called; default to V5 and re-resolve in start(). 

41 parser: MQTTBaseParser 

42 if getattr(config._outer_config, "version", "5.0") == "3.1.1": (empty)

43 parser = MQTTParserV311() (empty)

44 else: 

45 parser = MQTTParserV5() (empty)

46 config.parser = parser.parse_message (empty)

47 config.decoder = parser.decode_message (empty)

48 super().__init__(config, specification, calls) (empty)

49 self._topic = config.topic (empty)

50 self._shared = config.shared (empty)

51 self._qos = config.qos (empty)

52 self._subscription: zmqtt.Subscription | None = None (empty)

53 

54 if config.ack_policy is AckPolicy.NACK_ON_ERROR: 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true(empty)

55 warnings.warn( 

56 "MQTT has no nack primitive; with NACK_ON_ERROR, " 

57 "on error QoS 1/2 messages will not be acknowledged " 

58 "and the broker will redeliver them.", 

59 RuntimeWarning, 

60 stacklevel=3, 

61 ) 

62 

63 @property (empty)

64 def topic(self) -> str: (empty)

65 full = f"{self._outer_config.prefix}{self._topic}" (empty)

66 return f"$share/{self._shared}/{full}" if self._shared else full (empty)

67 

68 def _make_response_publisher( (empty)

69 self, 

70 message: "StreamMessage[Any]", 

71 ) -> Sequence["PublisherProto"]: 

72 return ( (empty)

73 MQTTFakePublisher( 

74 producer=self._outer_config.producer, 

75 topic=message.reply_to, 

76 ), 

77 ) 

78 

79 @staticmethod (empty)

80 def build_log_context( (empty)

81 message: "StreamMessage[zmqtt.Message] | None", 

82 topic: str = "", 

83 ) -> dict[str, str]: 

84 return { (empty)

85 "topic": topic, 

86 "message_id": getattr(message, "message_id", ""), 

87 } 

88 

89 def get_log_context( (empty)

90 self, 

91 message: "StreamMessage[zmqtt.Message] | None", 

92 ) -> dict[str, str]: 

93 return self.build_log_context(message=message, topic=self.topic) (empty)

94 

95 @override (empty)

96 async def start(self) -> None: (empty)

97 # Re-resolve the parser now that _outer_config is fully composed 

98 # (i.e. include_router has been called and the broker's MQTTBrokerConfig 

99 # is reachable through the config chain). 

100 parser: MQTTBaseParser 

101 if getattr(self._outer_config, "version", "5.0") == "3.1.1": (empty)

102 parser = MQTTParserV311() (empty)

103 else: 

104 parser = MQTTParserV5() (empty)

105 self._parser = parser.parse_message (empty)

106 self._decoder = parser.decode_message (empty)

107 

108 await super().start() (empty)

109 

110 if self.calls: (empty)

111 await self._create_subscription() (empty)

112 self.add_task(self._consume_loop) (empty)

113 

114 self._post_start() (empty)

115 

116 @override (empty)

117 async def stop(self) -> None: (empty)

118 await super().stop() (empty)

119 if self._subscription is not None: (empty)

120 with suppress(Exception): (empty)

121 await self._subscription.stop() (empty)

122 self._subscription = None (empty)

123 

124 async def _create_subscription(self) -> None: (empty)

125 auto_ack = self.ack_policy is AckPolicy.ACK_FIRST (empty)

126 self._subscription = self._outer_config.client.subscribe( (empty)

127 self.topic, 

128 qos=zmqtt.QoS(self._qos), 

129 auto_ack=auto_ack, 

130 ) 

131 await self._subscription.start() (empty)

132 

133 @override (empty)

134 async def get_one( (empty)

135 self, 

136 *, 

137 timeout: float = 5.0, 

138 ) -> "StreamMessage[zmqtt.Message] | None": 

139 assert not self.calls, ( (empty)

140 "You can't use `get_one` method if subscriber has registered handlers." 

141 ) 

142 

143 if self._subscription is None: 143 ↛ 152line 143 didn't jump to line 152 because the condition on line 143 was always true(empty)

144 auto_ack = self.ack_policy is AckPolicy.ACK_FIRST (empty)

145 self._subscription = self._outer_config.client.subscribe( (empty)

146 self.topic, 

147 qos=zmqtt.QoS(self._qos), 

148 auto_ack=auto_ack, 

149 ) 

150 await self._subscription.start() (empty)

151 

152 async_parser, async_decoder = self._get_parser_and_decoder() (empty)

153 

154 raw_msg: zmqtt.Message | None = None (empty)

155 with anyio.move_on_after(timeout): (empty)

156 raw_msg = await self._subscription.get_message() (empty)

157 

158 context = self._outer_config.fd_config.context (empty)

159 return await process_msg( (empty)

160 msg=raw_msg, 

161 middlewares=(m(raw_msg, context=context) for m in self._broker_middlewares), 

162 parser=async_parser, 

163 decoder=async_decoder, 

164 ) 

165 

166 @override (empty)

167 async def __aiter__(self) -> AsyncIterator["StreamMessage[zmqtt.Message]"]: # type: ignore[override] (empty)

168 if self._subscription is None: 168 ↛ 171line 168 didn't jump to line 171 because the condition on line 168 was always true(empty)

169 await self._create_subscription() (empty)

170 

171 assert self._subscription is not None (empty)

172 context = self._outer_config.fd_config.context (empty)

173 async_parser, async_decoder = self._get_parser_and_decoder() (empty)

174 async for raw_msg in self._subscription: 174 ↛ exitline 174 didn't return from function '__aiter__' because the loop on line 174 didn't complete(empty)

175 msg: MQTTMessage = await process_msg( # type: ignore[assignment] (empty)

176 msg=raw_msg, 

177 middlewares=( 

178 m(raw_msg, context=context) for m in self._broker_middlewares 

179 ), 

180 parser=async_parser, 

181 decoder=async_decoder, 

182 ) 

183 yield msg (empty)

184 

185 @abstractmethod (empty)

186 async def _consume_loop(self) -> None: (empty)

187 raise NotImplementedError 

188 

189 

190class MQTTDefaultSubscriber(MQTTBaseSubscriber): (empty)

191 """Sequential MQTT subscriber — processes one message at a time.""" 

192 

193 async def _consume_loop(self) -> None: (empty)

194 assert self._subscription is not None (empty)

195 async for msg in self._subscription: (empty)

196 await self.consume(msg) (empty)

197 

198 

199class MQTTConcurrentSubscriber(ConcurrentMixin[zmqtt.Message], MQTTBaseSubscriber): (empty)

200 """Concurrent MQTT subscriber — processes up to max_workers messages in parallel.""" 

201 

202 @override (empty)

203 async def start(self) -> None: (empty)

204 await super().start() 

205 self.start_consume_task() 

206 

207 async def _consume_loop(self) -> None: (empty)

208 assert self._subscription is not None 

209 async for msg in self._subscription: 

210 await self._put_msg(msg)