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

85 statements  

« prev     ^ index     » next       coverage.py v7.3.4, created at 2023-12-26 19:13 -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@typechecked 

168class RobotnikConfigTypedDict(TypedDict): 

169 tiers: List[List[Dict]] 

170 connection: Dict 

171 

172@typechecked 

173class RobotnikConfig(BaseModel): 

174 tiers: List[List[ServerConfig]] 

175 connection: ConnectionConfig = ConnectionConfig() 

176 

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

178 return self.tiers[index] 

179 

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

181 return choice(self.tier(tier)) 

182 

183 def as_dict(self) -> RobotnikConfigTypedDict: 

184 return self.dict() 

185 

186 @staticmethod 

187 def from_tiered( 

188 tiers: List[List[ServerConfig]], 

189 ) -> "RobotnikConfig": 

190 return RobotnikConfig(tiers=tiers) 

191 

192 @staticmethod 

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

194 return RobotnikConfig( 

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

196 ) 

197 

198 

199@typechecked 

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

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

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

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

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