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

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@typechecked

168 class RobotnikConfigTypedDict(TypedDict):

169 tiers: List[List[Dict]]

170 connection: Dict

171

172@typechecked

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

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