Coverage for zenmqtt/mqtt/protocol/context.py: 100%

1import asyncio 

2from typing import Optional 

3 

4from zenmqtt.connection import MQTTConnection 

5from zenmqtt.exceptions import NotConnectedError 

6from zenmqtt.instrumentation.metrics import MetricsCollector 

7from zenmqtt.mqtt.limits import DEFAULT_SERVER_LIMITS, ServerLimits 

8from zenmqtt.mqtt.packet import PacketType 

9 

10 

11class ProtocolContext: 

12 """ 

13 State shared by the parts of the protocol: the current connection, limits 

14 of the server and metrics. 

15 

16 The state is changed only by MQTTProtocol, with the methods below; the 

17 parts read it through the read-only properties. 

18 """ 

19 

20 def __init__(self, metrics: MetricsCollector) -> None: 

21 self.metrics = metrics 

22 

23 self._connection: Optional[MQTTConnection] = None 

24 self._connected = False 

25 self._server_limits: ServerLimits = DEFAULT_SERVER_LIMITS 

26 

27 @property 

28 def connection(self) -> Optional[MQTTConnection]: 

29 """The current connection, also before CONNACK and after its loss.""" 

30 return self._connection 

31 

32 @property 

33 def connected(self) -> bool: 

34 """True only between successful CONNACK and loss of the connection.""" 

35 return self._connected 

36 

37 @property 

38 def server_limits(self) -> ServerLimits: 

39 """Limits of the server from the last CONNACK.""" 

40 return self._server_limits 

41 

42 # called by MQTTProtocol only 

43 

44 def set_connection(self, connection: MQTTConnection) -> None: 

45 self._connection = connection 

46 

47 def set_server_limits(self, server_limits: ServerLimits) -> None: 

48 self._server_limits = server_limits 

49 

50 def mark_connected(self) -> None: 

51 self._connected = True 

52 

53 def mark_disconnected(self) -> None: 

54 self._connected = False 

55 

56 def ensure_connected(self) -> MQTTConnection: 

57 connection = self._connection 

58 

59 if not self._connected or not connection or connection.is_closing(): 

60 raise NotConnectedError() 

61 

62 return connection 

63 

64 async def write(self, connection: MQTTConnection, packet: bytes) -> None: 

65 await connection.write(packet) 

66 

67 self.metrics.on_packet_sent(PacketType(packet[0] >> 4), len(packet)) 

68 

69 @staticmethod 

70 def now() -> float: 

71 return asyncio.get_running_loop().time()