Coverage for faststream / _internal / broker / broker.py: 98%

1from abc import abstractmethod (empty)

2from collections.abc import Iterable, Sequence (empty)

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

4 

5from fast_depends import Provider (empty)

6from typing_extensions import Self (empty)

7 

8from faststream._internal.configs import BrokerConfigType (empty)

9from faststream._internal.types import ( (empty)

10 BrokerMiddleware, 

11 ConnectionType, 

12 MsgType, 

13) 

14 

15from .pub_base import BrokerPublishMixin (empty)

16from .registrator import Registrator (empty)

17 

18if TYPE_CHECKING: (empty)

19 from types import TracebackType 

20 

21 from faststream._internal.context.repository import ContextRepo 

22 from faststream._internal.di import FastDependsConfig 

23 from faststream._internal.producer import ProducerProto 

24 from faststream.specification.schema import BrokerSpec 

25 

26 

27class BrokerUsecase( (empty)

28 Registrator[MsgType, BrokerConfigType], 

29 BrokerPublishMixin[MsgType], 

30 Generic[MsgType, ConnectionType, BrokerConfigType], 

31): 

32 """Basic class for brokers-only. 

33 

34 Extends `Registrator` by connection, publish and AsyncAPI behavior. 

35 """ 

36 

37 _connection: ConnectionType | None (empty)

38 

39 def __init__( (empty)

40 self, 

41 *, 

42 config: BrokerConfigType, 

43 specification: "BrokerSpec", 

44 routers: Iterable[Registrator[Any, Any]], 

45 **connection_kwargs: Any, 

46 ) -> None: 

47 super().__init__( (empty)

48 routers=routers, 

49 config=config, 

50 ) 

51 self.specification = specification (empty)

52 

53 self.running = False (empty)

54 

55 self._connection_kwargs = connection_kwargs (empty)

56 self._connection = None (empty)

57 

58 @property (empty)

59 def middlewares(self) -> Sequence["BrokerMiddleware[MsgType]"]: (empty)

60 return self.config.broker_middlewares (empty)

61 

62 @property (empty)

63 def _producer(self) -> "ProducerProto": (empty)

64 return self.config.producer (empty)

65 

66 @property (empty)

67 def context(self) -> "ContextRepo": (empty)

68 return self.config.fd_config.context (empty)

69 

70 @property (empty)

71 def provider(self) -> Provider: (empty)

72 return self.config.fd_config.provider 

73 

74 async def __aenter__(self) -> "Self": (empty)

75 await self.connect() (empty)

76 return self (empty)

77 

78 async def __aexit__( (empty)

79 self, 

80 exc_type: type[BaseException] | None, 

81 exc_val: BaseException | None, 

82 exc_tb: Optional["TracebackType"], 

83 ) -> None: 

84 await self.stop(exc_type, exc_val, exc_tb) (empty)

85 

86 def _update_fd_config(self, config: "FastDependsConfig") -> None: (empty)

87 """Private method to change broker config state by outer application.""" 

88 self.config.fd_config = config | self.config.fd_config (empty)

89 

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

91 # TODO: filter by already running handlers after TestClient refactor 

92 for sub in self.subscribers: (empty)

93 await sub.start() (empty)

94 

95 for pub in self.publishers: (empty)

96 await pub.start() (empty)

97 

98 self.running = True (empty)

99 

100 def _setup_logger(self) -> None: (empty)

101 for sub in self.subscribers: (empty)

102 log_context = sub.get_log_context(None) (empty)

103 log_context.pop("message_id", None) (empty)

104 self.config.logger.params_storage.register_subscriber(log_context) (empty)

105 

106 self.config.logger._setup(self.config.fd_config.context) (empty)

107 

108 async def connect(self) -> ConnectionType: (empty)

109 """Connect to a remote server.""" 

110 if self._connection is None: (empty)

111 self._connection = await self._connect() (empty)

112 self._setup_logger() (empty)

113 

114 return self._connection (empty)

115 

116 @abstractmethod (empty)

117 async def _connect(self) -> ConnectionType: (empty)

118 raise NotImplementedError 

119 

120 async def stop( (empty)

121 self, 

122 exc_type: type[BaseException] | None = None, 

123 exc_val: BaseException | None = None, 

124 exc_tb: Optional["TracebackType"] = None, 

125 ) -> None: 

126 """Closes the object.""" 

127 for sub in self.subscribers: (empty)

128 await sub.stop() (empty)

129 

130 self.running = False (empty)

131 

132 @abstractmethod (empty)

133 async def ping(self, timeout: float | None) -> bool: (empty)

134 """Check connection alive.""" 

135 raise NotImplementedError