Coverage for faststream / middlewares / acknowledgement / middleware.py: 91%

1import logging (empty)

2from typing import TYPE_CHECKING, Any, Optional (empty)

3 

4from faststream._internal.middlewares import BaseMiddleware (empty)

5from faststream.exceptions import ( (empty)

6 AckMessage, 

7 HandlerException, 

8 NackMessage, 

9 RejectMessage, 

10) 

11 

12from .config import AckPolicy (empty)

13 

14if TYPE_CHECKING: (empty)

15 from types import TracebackType 

16 

17 from faststream._internal.basic_types import AsyncFuncAny 

18 from faststream._internal.context.repository import ContextRepo 

19 from faststream._internal.logger import LoggerState 

20 from faststream.message import StreamMessage 

21 

22 

23class AcknowledgementMiddleware: (empty)

24 def __init__( (empty)

25 self, 

26 logger: "LoggerState", 

27 ack_policy: "AckPolicy", 

28 extra_options: dict[str, Any], 

29 ) -> None: 

30 self.ack_policy = ack_policy (empty)

31 self.extra_options = extra_options (empty)

32 self.logger = logger (empty)

33 

34 def __call__( (empty)

35 self, 

36 msg: Any | None, 

37 context: "ContextRepo", 

38 ) -> "_AcknowledgementMiddleware": 

39 return _AcknowledgementMiddleware( (empty)

40 msg, 

41 logger=self.logger, 

42 ack_policy=self.ack_policy, 

43 extra_options=self.extra_options, 

44 context=context, 

45 ) 

46 

47 

48class _AcknowledgementMiddleware(BaseMiddleware): (empty)

49 def __init__( (empty)

50 self, 

51 msg: Any | None, 

52 /, 

53 *, 

54 logger: "LoggerState", 

55 context: "ContextRepo", 

56 extra_options: dict[str, Any], 

57 # can't be created with AckPolicy.MANUAL 

58 ack_policy: AckPolicy, 

59 ) -> None: 

60 super().__init__(msg, context=context) (empty)

61 

62 self.ack_policy = ack_policy (empty)

63 self.extra_options = extra_options (empty)

64 self.logger = logger (empty)

65 

66 self.message: StreamMessage[Any] | None = None (empty)

67 

68 async def consume_scope( (empty)

69 self, 

70 call_next: "AsyncFuncAny", 

71 msg: "StreamMessage[Any]", 

72 ) -> Any: 

73 self.message = msg (empty)

74 if self.ack_policy is AckPolicy.ACK_FIRST: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true(empty)

75 await self.__ack() 

76 

77 return await call_next(msg) (empty)

78 

79 async def __aexit__( (empty)

80 self, 

81 exc_type: type[BaseException] | None = None, 

82 exc_val: BaseException | None = None, 

83 exc_tb: Optional["TracebackType"] = None, 

84 ) -> bool | None: 

85 if self.ack_policy is AckPolicy.ACK_FIRST: 85 ↛ 86line 85 didn't jump to line 86 because the condition on line 85 was never true(empty)

86 return False 

87 

88 if not exc_type: (empty)

89 await self.__ack() (empty)

90 

91 elif isinstance(exc_val, HandlerException): (empty)

92 if isinstance(exc_val, AckMessage): (empty)

93 await self.__ack(**exc_val.extra_options) (empty)

94 

95 elif isinstance(exc_val, NackMessage): (empty)

96 await self.__nack(**exc_val.extra_options) (empty)

97 

98 elif isinstance(exc_val, RejectMessage): # pragma: no branch (empty)

99 await self.__reject(**exc_val.extra_options) (empty)

100 

101 # Exception was processed and suppressed 

102 return True (empty)

103 

104 elif self.ack_policy is AckPolicy.REJECT_ON_ERROR: (empty)

105 await self.__reject() (empty)

106 

107 elif self.ack_policy is AckPolicy.NACK_ON_ERROR: (empty)

108 await self.__nack() (empty)

109 

110 elif self.ack_policy is AckPolicy.ACK: 110 ↛ 114line 110 didn't jump to line 114 because the condition on line 110 was always true(empty)

111 await self.__ack() (empty)

112 

113 # Exception was not processed 

114 return False (empty)

115 

116 async def __ack(self, **exc_extra_options: Any) -> None: (empty)

117 if self.message: (empty)

118 try: (empty)

119 await self.message.ack(**exc_extra_options, **self.extra_options) (empty)

120 except Exception as er: (empty)

121 if self.logger is not None: (empty)

122 self.logger.log(repr(er), logging.CRITICAL, exc_info=er) (empty)

123 

124 async def __nack(self, **exc_extra_options: Any) -> None: (empty)

125 if self.message: 125 ↛ exitline 125 didn't return from function '__nack' because the condition on line 125 was always true(empty)

126 try: (empty)

127 await self.message.nack(**exc_extra_options, **self.extra_options) (empty)

128 except Exception as er: 

129 if self.logger is not None: 

130 self.logger.log(repr(er), logging.CRITICAL, exc_info=er) 

131 

132 async def __reject(self, **exc_extra_options: Any) -> None: (empty)

133 if self.message: (empty)

134 try: (empty)

135 await self.message.reject(**exc_extra_options, **self.extra_options) (empty)

136 except Exception as er: (empty)

137 if self.logger is not None: (empty)

138 self.logger.log(repr(er), logging.CRITICAL, exc_info=er) (empty)