Coverage for faststream / _internal / endpoint / call_wrapper.py: 96%

1import asyncio (empty)

2from collections.abc import Awaitable, Callable, Reversible, Sequence (empty)

3from typing import ( (empty)

4 TYPE_CHECKING, 

5 Any, 

6 Generic, 

7 Optional, 

8) 

9from unittest.mock import MagicMock (empty)

10 

11import anyio (empty)

12 

13from faststream._internal.types import P_HandlerParams, T_HandlerReturn (empty)

14from faststream.exceptions import SetupError (empty)

15 

16if TYPE_CHECKING: (empty)

17 from fast_depends.core import CallModel 

18 from fast_depends.dependencies import Dependant 

19 

20 from faststream._internal.basic_types import Decorator 

21 from faststream._internal.di import FastDependsConfig 

22 from faststream._internal.endpoint.publisher import PublisherProto 

23 from faststream._internal.endpoint.subscriber import SubscriberUsecase 

24 from faststream.message import StreamMessage 

25 

26 

27def ensure_call_wrapper( (empty)

28 call: Callable[P_HandlerParams, T_HandlerReturn], 

29) -> "HandlerCallWrapper[P_HandlerParams, T_HandlerReturn]": 

30 if isinstance(call, HandlerCallWrapper): (empty)

31 return call (empty)

32 

33 return HandlerCallWrapper(call) (empty)

34 

35 

36class HandlerCallWrapper(Generic[P_HandlerParams, T_HandlerReturn]): (empty)

37 """A generic class to wrap handler calls.""" 

38 

39 future: Optional["asyncio.Future[Any]"] (empty)

40 _wrapped_call: Callable[..., Awaitable[Any]] | None (empty)

41 _original_call: Callable[P_HandlerParams, T_HandlerReturn] (empty)

42 

43 _publishers: list["PublisherProto[Any]"] (empty)

44 

45 # we have to store subscribers here 

46 # to protect them from garbage collection 

47 _subscribers: list["SubscriberUsecase[Any]"] (empty)

48 

49 __slots__ = ( (empty)

50 "_original_call", 

51 "_publishers", 

52 "_subscribers", 

53 "_wrapped_call", 

54 "future", 

55 "is_test", 

56 "mock", 

57 ) 

58 

59 def __init__( (empty)

60 self, 

61 call: Callable[P_HandlerParams, T_HandlerReturn], 

62 ) -> None: 

63 """Initialize a handler.""" 

64 self._original_call = call (empty)

65 self._wrapped_call = None (empty)

66 

67 self._publishers = [] (empty)

68 self._subscribers = [] (empty)

69 

70 self.mock = MagicMock() (empty)

71 self.future = None (empty)

72 self.is_test = False (empty)

73 

74 def __call__( (empty)

75 self, 

76 *args: P_HandlerParams.args, 

77 **kwargs: P_HandlerParams.kwargs, 

78 ) -> T_HandlerReturn: 

79 """Calls the object as a function.""" 

80 return self._original_call(*args, **kwargs) (empty)

81 

82 async def call_wrapped( (empty)

83 self, 

84 message: "StreamMessage[Any]", 

85 ) -> Any: 

86 """Calls the wrapped function with the given message.""" 

87 assert self._wrapped_call, "You should use `set_wrapped` first" (empty)

88 if self.is_test: (empty)

89 self.mock(await message.decode()) (empty)

90 return await self._wrapped_call(message) (empty)

91 

92 def set_wrapped( (empty)

93 self, 

94 *, 

95 dependencies: Sequence["Dependant"], 

96 _call_decorators: Reversible["Decorator"], 

97 config: "FastDependsConfig", 

98 ) -> "CallModel": 

99 dependent = config.build_call( (empty)

100 self._original_call, 

101 dependencies=dependencies, 

102 call_decorators=_call_decorators, 

103 ) 

104 self._original_call = dependent.original_call (empty)

105 self._wrapped_call = dependent.wrapped_call (empty)

106 return dependent.dependent (empty)

107 

108 async def wait_call(self, timeout: float | None = None) -> None: (empty)

109 """Waits for a call with an optional timeout.""" 

110 assert self.future is not None, "You can use this method only with TestClient" (empty)

111 with anyio.fail_after(timeout): (empty)

112 await self.future (empty)

113 

114 def set_test(self) -> None: (empty)

115 self.is_test = True (empty)

116 self.mock.reset_mock() (empty)

117 self.refresh(with_mock=True) (empty)

118 

119 def reset_test(self) -> None: (empty)

120 self.is_test = False (empty)

121 self.mock.reset_mock() (empty)

122 self.future = None (empty)

123 

124 def trigger( (empty)

125 self, 

126 result: Any = None, 

127 error: BaseException | None = None, 

128 ) -> None: 

129 if not self.is_test: (empty)

130 return (empty)

131 

132 if self.future is None: 132 ↛ 133line 132 didn't jump to line 133 because the condition on line 132 was never true(empty)

133 msg = "You can use this method only with TestClient" 

134 raise SetupError(msg) 

135 

136 if self.future.done(): (empty)

137 self.future = asyncio.Future() (empty)

138 

139 if error: (empty)

140 self.future.set_exception(error) (empty)

141 

142 else: 

143 self.future.set_result(result) (empty)

144 

145 def refresh(self, with_mock: bool = False) -> None: (empty)

146 if asyncio.events._get_running_loop() is not None: (empty)

147 self.future = asyncio.Future() (empty)

148 

149 if with_mock and self.mock is not None: (empty)

150 self.mock.reset_mock() (empty)