Coverage for /Users/eugene/Development/robotnikmq/robotnikmq/config.py: 82%

85 statements  

« prev     ^ index     » next       coverage.py v7.3.4, created at 2023-12-26 23:29 -0500

1""" 

2Functions and objects related to the Configuration of RobotnikMQ 

3""" 

4from pathlib import Path 

5from random import choice 

6from ssl import create_default_context 

7from typing import Union, List, Optional, Dict, Any, TypedDict 

8 

9from pika import ConnectionParameters, SSLOptions # type: ignore 

10from pika.credentials import PlainCredentials # type: ignore 

11from pydantic import BaseModel # type: ignore 

12from pydantic import validator 

13from typeguard import typechecked 

14from yaml import safe_load # type: ignore 

15 

16from robotnikmq.error import ( 

17 NotConfigured, 

18 InvalidConfiguration, 

19) 

20from robotnikmq.log import log 

21 

22 

23@typechecked 

24def _existing_file_or_none(path: Union[str, Path, None]) -> Optional[Path]: 

25 """ 

26 Validates that a given path exists (either a string or Path object) and returns it or throws an exception. 

27 

28 Parameters: 

29 path (Union[str, Path]): Description 

30 

31 Raises: 

32 FileDoesNotExist: Description 

33 

34 Returns: 

35 Path: Validated path that exists as of when the function was run 

36 """ 

37 return Path(path).resolve(strict=True) if path is not None else None 

38 

39 

40class ServerConfig(BaseModel): 

41 """ 

42 Configuration object representing the configuration information required to connect to a single server 

43 """ 

44 

45 host: str 

46 port: int 

47 user: str 

48 password: str 

49 vhost: str 

50 ca_cert: Optional[Path] = None 

51 cert: Optional[Path] = None 

52 key: Optional[Path] = None 

53 _conn_params: Optional[ConnectionParameters] = None 

54 

55 _existing_ca_cert = validator("ca_cert", pre=True, always=True, allow_reuse=True)( 

56 _existing_file_or_none 

57 ) 

58 _existing_cert = validator("cert", pre=True, always=True, allow_reuse=True)( 

59 _existing_file_or_none 

60 ) 

61 _existing_key = validator("key", pre=True, always=True, allow_reuse=True)( 

62 _existing_file_or_none 

63 ) 

64 

65 class Config: 

66 json_encoders = { 

67 Path: str, 

68 } 

69 

70 @typechecked 

71 def conn_params(self) -> ConnectionParameters: 

72 if self._conn_params is not None: 72 ↛ 73line 72 didn't jump to line 73, because the condition on line 72 was never true

73 return self._conn_params 

74 if self.ca_cert is not None and self.cert is not None and self.key is not None: 74 ↛ 75line 74 didn't jump to line 75, because the condition on line 74 was never true

75 context = create_default_context(cafile=str(self.ca_cert)) 

76 context.load_cert_chain(self.cert, self.key) 

77 return ConnectionParameters( 

78 host=self.host, 

79 port=self.port, 

80 virtual_host=self.vhost, 

81 credentials=PlainCredentials(self.user, self.password), 

82 ssl_options=SSLOptions(context, self.host), 

83 ) 

84 context = create_default_context() 

85 return ConnectionParameters( 

86 host=self.host, 

87 port=self.port, 

88 virtual_host=self.vhost, 

89 credentials=PlainCredentials(self.user, self.password), 

90 ) 

91 

92 @typechecked 

93 @staticmethod 

94 def from_connection_params(conn_params: ConnectionParameters) -> "ServerConfig": 

95 return ServerConfig( 

96 host=conn_params.host, 

97 port=conn_params.port, 

98 user=getattr(conn_params.credentials, "username", ""), 

99 password=getattr(conn_params.credentials, "password", ""), 

100 vhost=conn_params.virtual_host, 

101 ) 

102 

103 

104@typechecked 

105def server_config( 

106 host: str, 

107 port: int, 

108 user: str, 

109 password: str, 

110 vhost: str, 

111 ca_cert: Union[str, Path, None] = None, 

112 cert: Union[str, Path, None] = None, 

113 key: Union[str, Path, None] = None, 

114) -> ServerConfig: 

115 """Generates a [`ServerConfig`][robotnikmq.config.ServerConfig] object while validating that the necessary certificate information. 

116 

117 Args: 

118 host (str): Description 

119 port (int): Description 

120 user (str): Description 

121 password (str): Description 

122 vhost (str): Description 

123 ca_cert (Union[str, Path]): Description 

124 cert (Union[str, Path]): Description 

125 key (Union[str, Path]): Description 

126 """ 

127 if ca_cert is not None and cert is not None and key is not None: 127 ↛ 128line 127 didn't jump to line 128, because the condition on line 127 was never true

128 ca_cert, cert, key = Path(ca_cert), Path(cert), Path(key) 

129 return ServerConfig( 

130 host=host, 

131 port=port, 

132 user=user, 

133 password=password, 

134 vhost=vhost, 

135 ca_cert=ca_cert, 

136 cert=cert, 

137 key=key, 

138 ) 

139 elif ca_cert is None and cert is None and key is None: 139 ↛ 140,   139 ↛ 1482 missed branches: 1) line 139 didn't jump to line 140, because the condition on line 139 was never true, 2) line 139 didn't jump to line 148, because the condition on line 139 was never false

140 return ServerConfig( 

141 host=host, 

142 port=port, 

143 user=user, 

144 password=password, 

145 vhost=vhost, 

146 ) 

147 else: 

148 raise InvalidConfiguration( 

149 "Either all public key encryption fields (cert, key, ca-cert) must be provided, or none of them." 

150 ) 

151 

152 

153class ConnectionConfig(BaseModel): 

154 attempts: int = 10 

155 wait_random_min_seconds: int = 2 

156 wait_random_max_seconds: int = 5 

157 

158 

159@typechecked 

160def conn_config(attempts: int, min_wait: int, max_wait: int) -> ConnectionConfig: 

161 return ConnectionConfig( 

162 attempts=attempts, 

163 wait_random_min_seconds=min_wait, 

164 wait_random_max_seconds=max_wait, 

165 ) 

166 

167 

168@typechecked 

169class RobotnikConfigTypedDict(TypedDict): 

170 tiers: List[List[Dict]] 

171 connection: Dict 

172 

173 

174@typechecked 

175class RobotnikConfig(BaseModel): 

176 tiers: List[List[ServerConfig]] 

177 connection: ConnectionConfig = ConnectionConfig() 

178 

179 def tier(self, index: int) -> List[ServerConfig]: 

180 return self.tiers[index] 

181 

182 def a_server(self, tier: int) -> ServerConfig: 

183 return choice(self.tier(tier)) 

184 

185 def as_dict(self) -> RobotnikConfigTypedDict: 

186 return self.model_dump() 

187 

188 @staticmethod 

189 def from_tiered( 

190 tiers: List[List[ServerConfig]], 

191 ) -> "RobotnikConfig": 

192 return RobotnikConfig(tiers=tiers) 

193 

194 @staticmethod 

195 def from_connection_params(conn_params: ConnectionParameters) -> "RobotnikConfig": 

196 return RobotnikConfig( 

197 tiers=[[ServerConfig.from_connection_params(conn_params)]] 

198 ) 

199 

200 

201@typechecked 

202def config_of(config_file: Optional[Path]) -> RobotnikConfig: 

203 if config_file is None or not config_file.exists(): 

204 log.critical("No valid RobotnikMQ configuration file was provided") 

205 raise NotConfigured("No valid RobotnikMQ configuration file was provided") 

206 return RobotnikConfig(**safe_load(config_file.open().read()))