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

1from time import time 

2from typing import Optional, Any, Dict, Union, List 

3from uuid import uuid4 as uuid 

4 

5from pika import BasicProperties 

6from typeguard import typechecked 

7 

8from robotnikmq.core import Robotnik, RobotnikConfig 

9from robotnikmq.rpc_server import RpcError, RpcResponse 

10from robotnikmq.utils import to_json 

11 

12 

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

20 

21 def _on_response(self, _, __, props: BasicProperties, body: bytes) -> None: 

22 self.response = body.decode() if self.corr_id == props.correlation_id else None 

23 

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