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

1 from collections import namedtuple

2 from random import uniform

3 from time import sleep

4 from traceback import format_exc

5 from typing import Optional, Callable, List, Generator

6

7 from pika.exceptions import AMQPError

8 from pika.exchange_type import ExchangeType

9 from typeguard import typechecked

10

11 from robotnikmq.config import RobotnikConfig

12 from robotnikmq.core import Robotnik, Message

13 from robotnikmq.error import MalformedMessage

14

15 OnMessageCallback = Callable[[Message], None]

16

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

18

19

20 class 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)