Coverage for faststream / nats / testing.py: 98%

1from collections.abc import Generator, Iterable, Iterator (empty)

2from contextlib import ExitStack, contextmanager (empty)

3from typing import TYPE_CHECKING, Any, Optional, cast (empty)

4from unittest.mock import AsyncMock (empty)

5 

6import anyio (empty)

7from nats.aio.msg import Msg (empty)

8from typing_extensions import override (empty)

9 

10from faststream._internal.endpoint.utils import ParserComposition (empty)

11from faststream._internal.testing.broker import TestBroker (empty)

12from faststream.exceptions import SubscriberNotFound (empty)

13from faststream.message import encode_message, gen_cor_id (empty)

14from faststream.nats.broker import NatsBroker (empty)

15from faststream.nats.parser import NatsParser (empty)

16from faststream.nats.publisher.producer import NatsFastProducer (empty)

17from faststream.nats.schemas.js_stream import is_subject_match_wildcard (empty)

18 

19if TYPE_CHECKING: (empty)

20 from fast_depends.library.serializer import SerializerProto 

21 

22 from faststream._internal.basic_types import SendableMessage 

23 from faststream._internal.configs.broker import ConfigComposition 

24 from faststream.nats.configs import NatsBrokerConfig 

25 from faststream.nats.publisher.usecase import LogicPublisher 

26 from faststream.nats.response import NatsPublishCommand 

27 from faststream.nats.subscriber.usecases.basic import LogicSubscriber 

28 

29__all__ = ("TestNatsBroker",) (empty)

30 

31 

32@contextmanager (empty)

33def change_producer( (empty)

34 config: "ConfigComposition[NatsBrokerConfig]", 

35 producer: "NatsFastProducer", 

36) -> Generator[None, None, None]: 

37 old_producer, config.broker_config.producer = ( (empty)

38 config.broker_config.producer, 

39 producer, 

40 ) 

41 old_js_producer, config.broker_config.js_producer = ( (empty)

42 config.broker_config.js_producer, 

43 producer, 

44 ) 

45 yield (empty)

46 config.broker_config.producer = old_producer (empty)

47 config.broker_config.js_producer = old_js_producer (empty)

48 

49 

50class TestNatsBroker(TestBroker[NatsBroker]): (empty)

51 """A class to test NATS brokers.""" 

52 

53 @staticmethod (empty)

54 def create_publisher_fake_subscriber( (empty)

55 broker: NatsBroker, 

56 publisher: "LogicPublisher", 

57 ) -> tuple["LogicSubscriber[Any]", bool]: 

58 publisher_stream = publisher.stream.name if publisher.stream else None (empty)

59 

60 sub: LogicSubscriber[Any] | None = None (empty)

61 for handler in broker.subscribers: (empty)

62 handler = cast("LogicSubscriber[Any]", handler) (empty)

63 if _is_handler_matches(handler, publisher.subject, publisher_stream): (empty)

64 sub = handler (empty)

65 break (empty)

66 

67 if sub is None: (empty)

68 is_real = False (empty)

69 sub = broker.subscriber( (empty)

70 publisher.subject, persistent=False, stream=publisher_stream 

71 ) 

72 else: 

73 is_real = True (empty)

74 

75 return sub, is_real (empty)

76 

77 @contextmanager (empty)

78 def _patch_producer(self, broker: NatsBroker) -> Iterator[None]: (empty)

79 fake_producer = FakeProducer(broker) (empty)

80 

81 with ExitStack() as es: (empty)

82 es.enter_context(change_producer(broker.config, fake_producer)) (empty)

83 yield (empty)

84 

85 async def _fake_connect( (empty)

86 self, 

87 broker: NatsBroker, 

88 *args: Any, 

89 **kwargs: Any, 

90 ) -> None: 

91 if not broker.config.connection_state: (empty)

92 broker.config.connection_state.connect(AsyncMock(), AsyncMock()) (empty)

93 

94 def _fake_start(self, broker: NatsBroker, *args: Any, **kwargs: Any) -> None: (empty)

95 if not broker.config.connection_state: (empty)

96 broker.config.connection_state.connect(AsyncMock(), AsyncMock()) (empty)

97 return super()._fake_start(broker, *args, **kwargs) (empty)

98 

99 

100class FakeProducer(NatsFastProducer): (empty)

101 def __init__(self, broker: NatsBroker) -> None: (empty)

102 self.broker = broker (empty)

103 

104 default = NatsParser(pattern="", is_ack_disabled=True) (empty)

105 self._parser = ParserComposition(broker._parser, default.parse_message) (empty)

106 self._decoder = ParserComposition(broker._decoder, default.decode_message) (empty)

107 

108 @override (empty)

109 async def publish(self, cmd: "NatsPublishCommand") -> None: (empty)

110 incoming = build_message( (empty)

111 message=cmd.body, 

112 subject=cmd.destination, 

113 headers=cmd.headers, 

114 correlation_id=cmd.correlation_id, 

115 reply_to=cmd.reply_to, 

116 serializer=self.broker.config.fd_config._serializer, 

117 ) 

118 

119 for handler in _find_handler( (empty)

120 cast("list[LogicSubscriber[Any]]", self.broker.subscribers), 

121 cmd.destination, 

122 cmd.stream, 

123 ): 

124 msg: list[PatchedMessage] | PatchedMessage 

125 

126 if (pull := getattr(handler, "pull_sub", None)) and pull.batch: (empty)

127 msg = [incoming] (empty)

128 else: 

129 msg = incoming (empty)

130 

131 await self._execute_handler(msg, cmd.destination, handler) (empty)

132 

133 @override (empty)

134 async def request(self, cmd: "NatsPublishCommand") -> "PatchedMessage": (empty)

135 incoming = build_message( (empty)

136 message=cmd.body, 

137 subject=cmd.destination, 

138 headers=cmd.headers, 

139 correlation_id=cmd.correlation_id, 

140 serializer=self.broker.config.fd_config._serializer, 

141 ) 

142 

143 for handler in _find_handler( (empty)

144 cast("list[LogicSubscriber[Any]]", self.broker.subscribers), 

145 cmd.destination, 

146 cmd.stream, 

147 ): 

148 msg: list[PatchedMessage] | PatchedMessage 

149 

150 if (pull := getattr(handler, "pull_sub", None)) and pull.batch: (empty)

151 msg = [incoming] (empty)

152 else: 

153 msg = incoming (empty)

154 

155 with anyio.fail_after(cmd.timeout): (empty)

156 return await self._execute_handler(msg, cmd.destination, handler) (empty)

157 

158 raise SubscriberNotFound (empty)

159 

160 async def _execute_handler( (empty)

161 self, 

162 msg: Any, 

163 subject: str, 

164 handler: "LogicSubscriber[Any]", 

165 ) -> "PatchedMessage": 

166 result = await handler.process_message(msg) (empty)

167 

168 return build_message( (empty)

169 subject=subject, 

170 message=result.body, 

171 headers=result.headers, 

172 correlation_id=result.correlation_id, 

173 serializer=self.broker.config.fd_config._serializer, 

174 ) 

175 

176 

177def _find_handler( (empty)

178 subscribers: Iterable["LogicSubscriber[Any]"], 

179 subject: str, 

180 stream: str | None = None, 

181) -> Generator["LogicSubscriber[Any]", None, None]: 

182 published_queues = set() (empty)

183 for handler in subscribers: (empty)

184 if _is_handler_matches(handler, subject, stream): (empty)

185 if queue := getattr(handler, "queue", None): (empty)

186 if queue in published_queues: (empty)

187 continue (empty)

188 else: 

189 published_queues.add(queue) (empty)

190 yield handler (empty)

191 

192 

193def _is_handler_matches( (empty)

194 handler: "LogicSubscriber[Any]", 

195 subject: str, 

196 stream: str | None = None, 

197) -> bool: 

198 if stream: (empty)

199 if not (handler_stream := getattr(handler, "stream", None)): (empty)

200 return False (empty)

201 

202 if stream != handler_stream.name: 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true(empty)

203 return False 

204 

205 if is_subject_match_wildcard(subject, handler.clear_subject): (empty)

206 return True (empty)

207 

208 for filter_subject in handler.filter_subjects or (): (empty)

209 if is_subject_match_wildcard(subject, filter_subject): (empty)

210 return True (empty)

211 

212 return False (empty)

213 

214 

215def build_message( (empty)

216 message: "SendableMessage", 

217 subject: str, 

218 *, 

219 reply_to: str = "", 

220 correlation_id: str | None = None, 

221 headers: dict[str, str] | None = None, 

222 serializer: Optional["SerializerProto"] = None, 

223) -> "PatchedMessage": 

224 msg, content_type = encode_message(message, serializer=serializer) (empty)

225 return PatchedMessage( (empty)

226 _client=None, # type: ignore[arg-type] 

227 subject=subject, 

228 reply=reply_to, 

229 data=msg, 

230 headers={ 

231 "content-type": content_type or "", 

232 "correlation_id": correlation_id or gen_cor_id(), 

233 **(headers or {}), 

234 }, 

235 ) 

236 

237 

238class PatchedMessage(Msg): (empty)

239 async def ack(self) -> None: (empty)

240 pass (empty)

241 

242 async def ack_sync( (empty)

243 self, 

244 timeout: float = 1, (empty)

245 ) -> "PatchedMessage": # pragma: no cover (empty)

246 return self 

247 

248 async def nak(self, delay: float | None = None) -> None: (empty)

249 pass 

250 

251 async def term(self) -> None: (empty)

252 pass 

253 

254 async def in_progress(self) -> None: (empty)

255 pass