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

1"""

2 Functions and objects related to the Configuration of RobotnikMQ

3"""

4 from pathlib import Path

5 from random import choice

6 from ssl import create_default_context

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

8

9 from pika import ConnectionParameters, SSLOptions # type: ignore

10 from pika.credentials import PlainCredentials # type: ignore

11 from pydantic import BaseModel # type: ignore

12 from pydantic import validator

13 from typeguard import typechecked

14 from yaml import safe_load # type: ignore

15

16 from robotnikmq.error import (

17 NotConfigured,

18 InvalidConfiguration,

19)

20 from robotnikmq.log import log

21

22

23@typechecked

24 def _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

40 class 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

105 def 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

153 class ConnectionConfig(BaseModel):

154 attempts: int = 10

155 wait_random_min_seconds: int = 2

156 wait_random_max_seconds: int = 5

157

158

159@typechecked

160 def 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

169 class RobotnikConfigTypedDict(TypedDict):

170 tiers: List[List[Dict]]

171 connection: Dict

172

173

174@typechecked

175 class 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

202 def 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()))