Coverage for /Users/eugene/Development/robotnikmq/robotnikmq/rpc_client.py: 30%

1 from time import time

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

3 from uuid import uuid4 as uuid

4

5 from pika import BasicProperties

6 from typeguard import typechecked

7

8 from robotnikmq.core import Robotnik, RobotnikConfig

9 from robotnikmq.rpc_server import RpcError, RpcResponse

10 from robotnikmq.utils import to_json

11

12

13@typechecked

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