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
« 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
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
16from robotnikmq.error import (
17 NotConfigured,
18 InvalidConfiguration,
19)
20from robotnikmq.log import log
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.
28 Parameters:
29 path (Union[str, Path]): Description
31 Raises:
32 FileDoesNotExist: Description
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
40class ServerConfig(BaseModel):
41 """
42 Configuration object representing the configuration information required to connect to a single server
43 """
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
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 )
65 class Config:
66 json_encoders = {
67 Path: str,
68 }
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 )
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 )
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.
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 )
153class ConnectionConfig(BaseModel):
154 attempts: int = 10
155 wait_random_min_seconds: int = 2
156 wait_random_max_seconds: int = 5
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 )
168@typechecked
169class RobotnikConfigTypedDict(TypedDict):
170 tiers: List[List[Dict]]
171 connection: Dict
174@typechecked
175class RobotnikConfig(BaseModel):
176 tiers: List[List[ServerConfig]]
177 connection: ConnectionConfig = ConnectionConfig()
179 def tier(self, index: int) -> List[ServerConfig]:
180 return self.tiers[index]
182 def a_server(self, tier: int) -> ServerConfig:
183 return choice(self.tier(tier))
185 def as_dict(self) -> RobotnikConfigTypedDict:
186 return self.model_dump()
188 @staticmethod
189 def from_tiered(
190 tiers: List[List[ServerConfig]],
191 ) -> "RobotnikConfig":
192 return RobotnikConfig(tiers=tiers)
194 @staticmethod
195 def from_connection_params(conn_params: ConnectionParameters) -> "RobotnikConfig":
196 return RobotnikConfig(
197 tiers=[[ServerConfig.from_connection_params(conn_params)]]
198 )
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()))