Merge branch 'adafruit:main' into connection-manager

Author: justmobilize

adafruit_minimqtt/adafruit_minimqtt.py

@@ -187,7 +187,7 @@ def __init__(
         self._is_connected = False
         self._msg_size_lim = MQTT_MSG_SZ_LIM
         self._pid = 0
-        self._timestamp: float = 0
+        self._last_msg_sent_timestamp: float = 0
         self.logger = NullLogger()
         """An optional logging attribute that can be set with with a Logger
         to enable debug logging."""
@@ -537,6 +537,7 @@ def _connect(
         if self._username is not None:
             self._send_str(self._username)
             self._send_str(self._password)
+        self._last_msg_sent_timestamp = self.get_monotonic_time()
         self.logger.debug("Receiving CONNACK packet from broker")
         stamp = self.get_monotonic_time()
         while True:
@@ -591,6 +592,7 @@ def disconnect(self) -> None:
         self._connection_manager.free_socket(self._sock)
         self._is_connected = False
         self._subscribed_topics = []
+        self._last_msg_sent_timestamp = 0
         if self.on_disconnect is not None:
             self.on_disconnect(self, self.user_data, 0)
 
@@ -604,6 +606,7 @@ def ping(self) -> list[int]:
         self._sock.send(MQTT_PINGREQ)
         ping_timeout = self.keep_alive
         stamp = self.get_monotonic_time()
+        self._last_msg_sent_timestamp = stamp
         rc, rcs = None, []
         while rc != MQTT_PINGRESP:
             rc = self._wait_for_msg()
@@ -678,6 +681,7 @@ def publish(
         self._sock.send(pub_hdr_fixed)
         self._sock.send(pub_hdr_var)
         self._sock.send(msg)
+        self._last_msg_sent_timestamp = self.get_monotonic_time()
         if qos == 0 and self.on_publish is not None:
             self.on_publish(self, self.user_data, topic, self._pid)
         if qos == 1:
@@ -755,6 +759,7 @@ def subscribe(self, topic: Optional[Union[tuple, str, list]], qos: int = 0) -> N
         self.logger.debug(f"payload: {payload}")
         self._sock.send(payload)
         stamp = self.get_monotonic_time()
+        self._last_msg_sent_timestamp = stamp
         while True:
             op = self._wait_for_msg()
             if op is None:
@@ -830,6 +835,7 @@ def unsubscribe(self, topic: Optional[Union[str, list]]) -> None:
         for t in topics:
             self.logger.debug(f"UNSUBSCRIBING from topic {t}")
         self._sock.send(payload)
+        self._last_msg_sent_timestamp = self.get_monotonic_time()
         self.logger.debug("Waiting for UNSUBACK...")
         while True:
             stamp = self.get_monotonic_time()
@@ -919,31 +925,41 @@ def reconnect(self, resub_topics: bool = True) -> int:
         return ret
 
     def loop(self, timeout: float = 0) -> Optional[list[int]]:
-        # pylint: disable = too-many-return-statements
         """Non-blocking message loop. Use this method to check for incoming messages.
         Returns list of packet types of any messages received or None.
 
         :param float timeout: return after this timeout, in seconds.
 
         """
+        if timeout < self._socket_timeout:
+            raise MMQTTException(
+                # pylint: disable=consider-using-f-string
+                "loop timeout ({}) must be bigger ".format(timeout)
+                + "than socket timeout ({}))".format(self._socket_timeout)
+            )
+
         self._connected()
         self.logger.debug(f"waiting for messages for {timeout} seconds")
-        if self._timestamp == 0:
-            self._timestamp = self.get_monotonic_time()
-        current_time = self.get_monotonic_time()
-        if current_time - self._timestamp >= self.keep_alive:
-            self._timestamp = 0
-            # Handle KeepAlive by expecting a PINGREQ/PINGRESP from the server
-            self.logger.debug(
-                "KeepAlive period elapsed - requesting a PINGRESP from the server..."
-            )
-            rcs = self.ping()
-            return rcs
 
         stamp = self.get_monotonic_time()
         rcs = []
 
         while True:
+            if (
+                self.get_monotonic_time() - self._last_msg_sent_timestamp
+                >= self.keep_alive
+            ):
+                # Handle KeepAlive by expecting a PINGREQ/PINGRESP from the server
+                self.logger.debug(
+                    "KeepAlive period elapsed - requesting a PINGRESP from the server..."
+                )
+                rcs.extend(self.ping())
+                # ping() itself contains a _wait_for_msg() loop which might have taken a while,
+                # so check here as well.
+                if self.get_monotonic_time() - stamp > timeout:
+                    self.logger.debug(f"Loop timed out after {timeout} seconds")
+                    break
+
             rc = self._wait_for_msg()
             if rc is not None:
                 rcs.append(rc)
@@ -953,11 +969,13 @@ def loop(self, timeout: float = 0) -> Optional[list[int]]:
 
         return rcs if rcs else None
 
-    def _wait_for_msg(self) -> Optional[int]:
+    def _wait_for_msg(self, timeout: Optional[float] = None) -> Optional[int]:
         # pylint: disable = too-many-return-statements
 
         """Reads and processes network events.
         Return the packet type or None if there is nothing to be received.
+
+        :param float timeout: return after this timeout, in seconds.
         """
         # CPython socket module contains a timeout attribute
         if hasattr(self._socket_pool, "timeout"):
@@ -967,7 +985,7 @@ def _wait_for_msg(self) -> Optional[int]:
                 return None
         else:  # socketpool, esp32spi
             try:
-                res = self._sock_exact_recv(1)
+                res = self._sock_exact_recv(1, timeout=timeout)
             except OSError as error:
                 if error.errno in (errno.ETIMEDOUT, errno.EAGAIN):
                     # raised by a socket timeout if 0 bytes were present
@@ -1036,7 +1054,9 @@ def _decode_remaining_length(self) -> int:
                 return n
             sh += 7
 
-    def _sock_exact_recv(self, bufsize: int) -> bytearray:
+    def _sock_exact_recv(
+        self, bufsize: int, timeout: Optional[float] = None
+    ) -> bytearray:
         """Reads _exact_ number of bytes from the connected socket. Will only return
         bytearray with the exact number of bytes requested.
 
@@ -1047,6 +1067,7 @@ def _sock_exact_recv(self, bufsize: int) -> bytearray:
         bytes is returned or trigger a timeout exception.
 
         :param int bufsize: number of bytes to receive
+        :param float timeout: timeout, in seconds. Defaults to keep_alive
         :return: byte array
         """
         stamp = self.get_monotonic_time()
@@ -1058,7 +1079,7 @@ def _sock_exact_recv(self, bufsize: int) -> bytearray:
             to_read = bufsize - recv_len
             if to_read < 0:
                 raise MMQTTException(f"negative number of bytes to read: {to_read}")
-            read_timeout = self.keep_alive
+            read_timeout = timeout if timeout is not None else self.keep_alive
             mv = mv[recv_len:]
             while to_read > 0:
                 recv_len = self._sock.recv_into(mv, to_read)

tests/test_loop.py

@@ -8,22 +8,109 @@
 import socket
 import ssl
 import time
+import errno
+
 from unittest import TestCase, main
 from unittest.mock import patch
+from unittest import mock
 
 import adafruit_minimqtt.adafruit_minimqtt as MQTT
 
 
+class Nulltet:
+    """
+    Mock Socket that does nothing.
+
+    Inspired by the Mocket class from Adafruit_CircuitPython_Requests
+    """
+
+    def __init__(self):
+        self.sent = bytearray()
+
+        self.timeout = mock.Mock()
+        self.connect = mock.Mock()
+        self.close = mock.Mock()
+
+    def send(self, bytes_to_send):
+        """
+        Record the bytes. return the length of this bytearray.
+        """
+        self.sent.extend(bytes_to_send)
+        return len(bytes_to_send)
+
+    # MiniMQTT checks for the presence of "recv_into" and switches behavior based on that.
+    # pylint: disable=unused-argument,no-self-use
+    def recv_into(self, retbuf, bufsize):
+        """Always raise timeout exception."""
+        exc = OSError()
+        exc.errno = errno.ETIMEDOUT
+        raise exc
+
+
+class Pingtet:
+    """
+    Mock Socket tailored for PINGREQ testing.
+    Records sent data, hands out PINGRESP for each PINGREQ received.
+
+    Inspired by the Mocket class from Adafruit_CircuitPython_Requests
+    """
+
+    PINGRESP = bytearray([0xD0, 0x00])
+
+    def __init__(self):
+        self._to_send = self.PINGRESP
+
+        self.sent = bytearray()
+
+        self.timeout = mock.Mock()
+        self.connect = mock.Mock()
+        self.close = mock.Mock()
+
+        self._got_pingreq = False
+
+    def send(self, bytes_to_send):
+        """
+        Recognize PINGREQ and record the indication that it was received.
+        Assumes it was sent in one chunk (of 2 bytes).
+        Also record the bytes. return the length of this bytearray.
+        """
+        self.sent.extend(bytes_to_send)
+        if bytes_to_send == b"\xc0\0":
+            self._got_pingreq = True
+        return len(bytes_to_send)
+
+    # MiniMQTT checks for the presence of "recv_into" and switches behavior based on that.
+    def recv_into(self, retbuf, bufsize):
+        """
+        If the PINGREQ indication is on, return PINGRESP, otherwise raise timeout exception.
+        """
+        if self._got_pingreq:
+            size = min(bufsize, len(self._to_send))
+            if size == 0:
+                return size
+            chop = self._to_send[0:size]
+            retbuf[0:] = chop
+            self._to_send = self._to_send[size:]
+            if len(self._to_send) == 0:
+                self._got_pingreq = False
+                self._to_send = self.PINGRESP
+            return size
+
+        exc = OSError()
+        exc.errno = errno.ETIMEDOUT
+        raise exc
+
+
 class Loop(TestCase):
     """basic loop() test"""
 
     connect_times = []
     INITIAL_RCS_VAL = 42
     rcs_val = INITIAL_RCS_VAL
 
-    def fake_wait_for_msg(self):
+    def fake_wait_for_msg(self, timeout=1):
         """_wait_for_msg() replacement. Sleeps for 1 second and returns an integer."""
-        time.sleep(1)
+        time.sleep(timeout)
         retval = self.rcs_val
         self.rcs_val += 1
         return retval
@@ -54,6 +141,8 @@ def test_loop_basic(self) -> None:
 
             time_before = time.monotonic()
             timeout = random.randint(3, 8)
+            # pylint: disable=protected-access
+            mqtt_client._last_msg_sent_timestamp = mqtt_client.get_monotonic_time()
             rcs = mqtt_client.loop(timeout=timeout)
             time_after = time.monotonic()
 
@@ -62,12 +151,33 @@ def test_loop_basic(self) -> None:
 
             # Check the return value.
             assert rcs is not None
-            assert len(rcs) > 1
+            assert len(rcs) >= 1
             expected_rc = self.INITIAL_RCS_VAL
+            # pylint: disable=not-an-iterable
             for ret_code in rcs:
                 assert ret_code == expected_rc
                 expected_rc += 1
 
+    # pylint: disable=invalid-name
+    def test_loop_timeout_vs_socket_timeout(self):
+        """
+        loop() should throw MMQTTException if the timeout argument
+        is bigger than the socket timeout.
+        """
+        mqtt_client = MQTT.MQTT(
+            broker="127.0.0.1",
+            port=1883,
+            socket_pool=socket,
+            ssl_context=ssl.create_default_context(),
+            socket_timeout=1,
+        )
+
+        mqtt_client.is_connected = lambda: True
+        with self.assertRaises(MQTT.MMQTTException) as context:
+            mqtt_client.loop(timeout=0.5)
+
+        assert "loop timeout" in str(context.exception)
+
     def test_loop_is_connected(self):
         """
         loop() should throw MMQTTException if not connected
@@ -84,6 +194,71 @@ def test_loop_is_connected(self):
 
         assert "not connected" in str(context.exception)
 
+    # pylint: disable=no-self-use
+    def test_loop_ping_timeout(self):
+        """Verify that ping will be sent even with loop timeout bigger than keep alive timeout
+        and no outgoing messages are sent."""
+
+        recv_timeout = 2
+        keep_alive_timeout = recv_timeout * 2
+        mqtt_client = MQTT.MQTT(
+            broker="localhost",
+            port=1883,
+            ssl_context=ssl.create_default_context(),
+            connect_retries=1,
+            socket_timeout=1,
+            recv_timeout=recv_timeout,
+            keep_alive=keep_alive_timeout,
+        )
+
+        # patch is_connected() to avoid CONNECT/CONNACK handling.
+        mqtt_client.is_connected = lambda: True
+        mocket = Pingtet()
+        # pylint: disable=protected-access
+        mqtt_client._sock = mocket
+
+        start = time.monotonic()
+        res = mqtt_client.loop(timeout=2 * keep_alive_timeout)
+        assert time.monotonic() - start >= 2 * keep_alive_timeout
+        assert len(mocket.sent) > 0
+        assert len(res) == 2
+        assert set(res) == {int(0xD0)}
+
+    # pylint: disable=no-self-use
+    def test_loop_ping_vs_msgs_sent(self):
+        """Verify that ping will not be sent unnecessarily."""
+
+        recv_timeout = 2
+        keep_alive_timeout = recv_timeout * 2
+        mqtt_client = MQTT.MQTT(
+            broker="localhost",
+            port=1883,
+            ssl_context=ssl.create_default_context(),
+            connect_retries=1,
+            socket_timeout=1,
+            recv_timeout=recv_timeout,
+            keep_alive=keep_alive_timeout,
+        )
+
+        # patch is_connected() to avoid CONNECT/CONNACK handling.
+        mqtt_client.is_connected = lambda: True
+
+        # With QoS=0 no PUBACK message is sent, so Nulltet can be used.
+        mocket = Nulltet()
+        # pylint: disable=protected-access
+        mqtt_client._sock = mocket
+
+        i = 0
+        topic = "foo"
+        message = "bar"
+        for _ in range(3 * keep_alive_timeout):
+            mqtt_client.publish(topic, message, qos=0)
+            mqtt_client.loop(1)
+            i += 1
+
+        # This means no other messages than the PUBLISH messages generated by the code above.
+        assert len(mocket.sent) == i * (2 + 2 + len(topic) + len(message))
+
 
 if __name__ == "__main__":
     main()