Coverage for /Users/eugene/Development/robotnikmq/robotnikmq/subscriber.py: 98%

67 statements  

« prev     ^ index     » next       coverage.py v7.3.4, created at 2023-12-26 23:29 -0500

1from collections import namedtuple 

2from random import uniform 

3from time import sleep 

4from traceback import format_exc 

5from typing import Optional, Callable, List, Generator 

6 

7from pika.exceptions import AMQPError 

8from pika.exchange_type import ExchangeType 

9from typeguard import typechecked 

10 

11from robotnikmq.config import RobotnikConfig 

12from robotnikmq.core import Robotnik, Message 

13from robotnikmq.error import MalformedMessage 

14 

15OnMessageCallback = Callable[[Message], None] 

16 

17ExchangeBinding = namedtuple("ExchangeBinding", ["exchange", "binding_key"]) 

18 

19 

20class Subscriber(Robotnik): 

21 TIMEOUT_MAX = 30 

22 TIMEOUT_MIN = 0.5 

23 TIMEOUT_STEP = 3 

24 TIMEOUT_JITTER = 2 

25 

26 @typechecked 

27 def __init__(self, exchange_bindings: Optional[List[ExchangeBinding]] = None, 

28 config: Optional[RobotnikConfig] = None,): 

29 super().__init__(config=config) 

30 self.exchange_bindings = exchange_bindings or [] 

31 

32 @typechecked 

33 def _bind(self, exchange_binding: ExchangeBinding) -> "Subscriber": 

34 self.exchange_bindings.append(exchange_binding) 

35 return self 

36 

37 @typechecked 

38 def bind(self, exchange: str, binding_key: str = "#") -> "Subscriber": 

39 return self._bind(ExchangeBinding(exchange, binding_key)) 

40 

41 @typechecked 

42 def _consume( 

43 self, inactivity_timeout: Optional[float] 

44 ) -> Generator[Optional[Message], None, None]: 

45 with self.open_channel() as channel: 45 ↛ exitline 45 didn't return from function '_consume'

46 queue_name = ( 

47 channel.queue_declare(queue="", exclusive=True).method.queue or "" 

48 ) 

49 for ex_b in self.exchange_bindings: 

50 channel.exchange_declare( 

51 exchange=ex_b.exchange, 

52 exchange_type=ExchangeType.topic, 

53 auto_delete=True, 

54 ) 

55 channel.queue_bind( 

56 exchange=ex_b.exchange, 

57 queue=queue_name, 

58 routing_key=ex_b.binding_key, 

59 ) 

60 try: 

61 for method, ___, body in channel.consume( 

62 queue=queue_name, # pragma: no cover 

63 auto_ack=False, 

64 inactivity_timeout=inactivity_timeout, 

65 ): 

66 if method and body: 

67 channel.basic_ack(delivery_tag=method.delivery_tag or 0) 

68 try: 

69 yield Message.of_json(body.decode()) 

70 except MalformedMessage: 

71 self.log.debug(format_exc()) 

72 else: 

73 yield None 

74 finally: 

75 try: 

76 channel.cancel() 

77 self.close_channel(channel) 

78 except AssertionError as exc: 

79 self.log.warning(f"Unable to close channel: {exc}") 

80 

81 @typechecked 

82 @staticmethod 

83 def _jitter(step: float, jitter: float) -> float: 

84 return uniform(step - jitter, step + jitter) 

85 

86 @typechecked 

87 def _backoff_with_jitter(self, current_timeout: float, 

88 timeout_min: Optional[float] = None, 

89 timeout_step: Optional[float] = None, 

90 timeout_jitter: Optional[float] = None, 

91 timeout_max: Optional[float] = None) -> float: 

92 timeout_min = timeout_min or self.TIMEOUT_MIN 

93 timeout_step = timeout_step or self.TIMEOUT_STEP 

94 timeout_jitter = timeout_jitter or self.TIMEOUT_JITTER 

95 timeout_max = timeout_max or self.TIMEOUT_MAX 

96 self.log.warning(f"Backing off for {current_timeout} seconds before reconnecting...") 

97 sleep(current_timeout) 

98 self.log.warning("Reconnecting") 

99 return min(current_timeout + self._jitter(timeout_step, timeout_jitter), timeout_max) 

100 

101 @typechecked 

102 def consume( 

103 self, inactivity_timeout: Optional[float] = None 

104 ) -> Generator[Optional[Message], None, None]: 

105 timeout = Subscriber.TIMEOUT_MIN 

106 while 42: 

107 try: 

108 for msg in self._consume(inactivity_timeout): 108 ↛ 106line 108 didn't jump to line 106, because the loop on line 108 didn't complete

109 yield msg 

110 timeout = Subscriber.TIMEOUT_MIN 

111 except (AMQPError, OSError) as exc: 

112 self.log.warning(f"Subscriber issue encountered: {exc}") 

113 timeout = self._backoff_with_jitter(timeout)