Coverage for /Users/eugene/Development/robotnikmq/robotnikmq/rpc_client.py: 30%
38 statements
« prev ^ index » next coverage.py v7.3.4, created at 2023-12-26 19:16 -0500
« prev ^ index » next coverage.py v7.3.4, created at 2023-12-26 19:16 -0500
1from time import time
2from typing import Optional, Any, Dict, Union, List
3from uuid import uuid4 as uuid
5from pika import BasicProperties
6from typeguard import typechecked
8from robotnikmq.core import Robotnik, RobotnikConfig
9from robotnikmq.rpc_server import RpcError, RpcResponse
10from robotnikmq.utils import to_json
13@typechecked
14class RpcClient(Robotnik):
15 def __init__(self, config: Optional[RobotnikConfig] = None):
16 super().__init__(config=config)
17 self.response: Optional[Any] = None
18 self.callback_queue = None
19 self.corr_id = str(uuid())
21 def _on_response(self, _, __, props: BasicProperties, body: bytes) -> None:
22 self.response = body.decode() if self.corr_id == props.correlation_id else None
24 def call(self, queue: str,
25 args: Optional[Dict[str, Any]] = None,
26 str_args: Optional[str] = None,
27 timeout: Optional[float] = None,
28 raise_on_error: bool = False) -> Union[RpcError, str, int, float,
29 Dict, List[Dict]]:
30 with self.open_channel() as channel:
31 result = channel.queue_declare(queue='', exclusive=True)
32 self.callback_queue = result.method.queue
33 channel.basic_consume(
34 queue=self.callback_queue,
35 on_message_callback=self._on_response,
36 auto_ack=True)
37 self.response = None
38 self.corr_id = str(uuid())
39 str_args = str_args or '{}'
40 channel.basic_publish(exchange='',
41 routing_key=queue,
42 properties=BasicProperties(reply_to=self.callback_queue,
43 correlation_id=self.corr_id),
44 body=to_json(args) if args is not None else str_args)
45 start_time = time()
46 while self.response is None:
47 self.connection.process_data_events()
48 if timeout is not None and time() > start_time + timeout:
49 raise TimeoutError(f'No response has been received for Request: {self.corr_id}')
50 response = RpcResponse.from_json(self.response)
51 if response is not None:
52 return response.data
53 error = RpcError.from_json(self.response)
54 if raise_on_error:
55 raise RuntimeError(error.details)
56 return error