Merge pull request #94 from etiennedm/smpclient-timeout-refactor

Author: JPHutchins

smpclient/__init__.py

@@ -63,6 +63,7 @@ class SMPClient:
     Args:
         transport: the `SMPTransport` to use
         address: the address of the SMP server, see `smpclient.transport` for details
+        timeout_s: the default timeout in seconds for SMP requests
 
     Example:
 
@@ -86,25 +87,30 @@ async def main():
     ```
     """
 
-    def __init__(self, transport: SMPTransport, address: str):  # noqa: DOC301
+    def __init__(
+        self, transport: SMPTransport, address: str, timeout_s: float = 2.5
+    ):  # noqa: DOC301
         self._transport: Final = transport
         self._address: Final = address
+        self._timeout_s = timeout_s
 
-    async def connect(self, timeout_s: float = 5.0) -> None:
+    async def connect(self, connect_timeout_s: float | None = None) -> None:
         """Connect to the SMP server.
 
         Args:
-            timeout_s: the timeout for the connection attempt in seconds
+            connect_timeout_s: the timeout for the connection attempt in seconds
         """
-        await self._transport.connect(self._address, timeout_s)
-        await self._initialize()
+        connect_timeout_s = connect_timeout_s if connect_timeout_s is not None else self._timeout_s
+
+        await self._transport.connect(self._address, connect_timeout_s)
+        await self._initialize(self._timeout_s)
 
     async def disconnect(self) -> None:
         """Disconnect from the SMP server."""
         await self._transport.disconnect()
 
     async def request(
-        self, request: SMPRequest[TRep, TEr1, TEr2], timeout_s: float = 120.000
+        self, request: SMPRequest[TRep, TEr1, TEr2], timeout_s: float | None = None
     ) -> TRep | TEr1 | TEr2:
         """Make an `SMPRequest` to the SMP server and return the Response or Error.
 
@@ -159,6 +165,7 @@ async def request(
         ```
 
         """
+        timeout_s = timeout_s if timeout_s is not None else self._timeout_s
 
         try:
             async with timeout(timeout_s):
@@ -198,8 +205,8 @@ async def upload(
         image: bytes,
         slot: int = 0,
         upgrade: bool = False,
-        first_timeout_s: float = 40.000,
-        subsequent_timeout_s: float = 2.500,
+        first_timeout_s: float = 40.0,
+        subsequent_timeout_s: float | None = None,
         use_sha: bool = True,
     ) -> AsyncIterator[int]:
         """Iteratively upload an `image` to `slot`, yielding the offset.
@@ -214,6 +221,8 @@ async def upload(
                 [boot_write_img_confirmed()](https://docs.zephyrproject.org/apidoc/latest/group__mcuboot__api.html#ga95ccc9e1c7460fec16b9ce9ac8ad7a72)
                 for this purpose.
             first_timeout_s: the timeout for the first `ImageUploadWrite` request
+                which might take longer than subsequent requests (e.g. if a big
+                chunk of flash memory has to be erased upfront).
             subsequent_timeout_s: the timeout for subsequent `ImageUploadWrite` requests
             use_sha: `True` to include the SHA256 hash of the image in the first
                 packet.
@@ -228,6 +237,9 @@ async def upload(
         Raises:
             SMPUploadError: if the upload routine fails
         """
+        subsequent_timeout_s = (
+            subsequent_timeout_s if subsequent_timeout_s is not None else self._timeout_s
+        )
 
         response = await self.request(
             self._maximize_image_upload_write_packet(
@@ -290,7 +302,7 @@ async def upload_file(
         self,
         file_data: bytes,
         file_path: str,
-        timeout_s: float = 2.500,
+        timeout_s: float | None = None,
     ) -> AsyncIterator[int]:
         """Iteratively upload a `file_data` to `file_path`, yielding the offset.
 
@@ -305,6 +317,8 @@ async def upload_file(
         Raises:
             SMPUploadError: if the upload routine fails
         """
+        timeout_s = timeout_s if timeout_s is not None else self._timeout_s
+
         response = await self.request(
             self._maximize_file_upload_packet(
                 FileUpload(name=file_path, off=0, data=b"", len=len(file_data)),
@@ -342,7 +356,7 @@ async def upload_file(
     async def download_file(
         self,
         file_path: str,
-        timeout_s: float = 2.500,
+        timeout_s: float | None = None,
     ) -> bytes:
         """Download a file from the SMP server.
 
@@ -356,6 +370,8 @@ async def download_file(
         Raises:
             SMPUploadError: if the download routine fails
         """
+        timeout_s = timeout_s if timeout_s is not None else self._timeout_s
+
         response = await self.request(FileDownload(off=0, name=file_path), timeout_s=timeout_s)
         file_length = 0
 
@@ -500,18 +516,18 @@ def _maximize_file_upload_packet(self, request: FileUpload, data: bytes) -> File
             len=request.len,
         )
 
-    async def _initialize(self) -> None:
+    async def _initialize(self, timeout_s: float | None = None) -> None:
         """Gather initialization information from the SMP server."""
+        timeout_s = timeout_s if timeout_s is not None else self._timeout_s
 
         try:
-            async with timeout(2):
-                mcumgr_parameters = await self.request(MCUMgrParametersRead())
-                if success(mcumgr_parameters):
-                    logger.debug(f"MCUMgr parameters: {mcumgr_parameters}")
-                    self._transport.initialize(mcumgr_parameters.buf_size)
-                elif error(mcumgr_parameters):
-                    logger.warning(f"Error reading MCUMgr parameters: {mcumgr_parameters}")
-                else:
-                    assert_never(mcumgr_parameters)
-        except asyncio.TimeoutError:
+            mcumgr_parameters = await self.request(MCUMgrParametersRead(), timeout_s=timeout_s)
+            if success(mcumgr_parameters):
+                logger.debug(f"MCUMgr parameters: {mcumgr_parameters}")
+                self._transport.initialize(mcumgr_parameters.buf_size)
+            elif error(mcumgr_parameters):
+                logger.warning(f"Error reading MCUMgr parameters: {mcumgr_parameters}")
+            else:
+                assert_never(mcumgr_parameters)
+        except TimeoutError:
             logger.warning("Timeout waiting for MCUMgr parameters")

tests/test_smp_client.py

@@ -85,12 +85,12 @@ def test_constructor() -> None:
 @pytest.mark.asyncio
 async def test_connect() -> None:
     m = SMPMockTransport()
-    s = SMPClient(m, "address")
+    s = SMPClient(m, "address", 5.0)
     s._initialize = AsyncMock()  # type: ignore
     await s.connect()
 
     m.connect.assert_awaited_once_with("address", 5.0)
-    s._initialize.assert_awaited_once_with()
+    s._initialize.assert_awaited_once_with(5.0)
 
 
 @pytest.mark.asyncio
@@ -175,7 +175,7 @@ async def test_request() -> None:
 @pytest.mark.asyncio
 async def test_upload() -> None:
     m = SMPMockTransport()
-    s = SMPClient(m, "address")
+    s = SMPClient(m, "address", 2.5)
 
     s.request = AsyncMock()  # type: ignore
 
@@ -239,7 +239,7 @@ async def test_upload() -> None:
             off=415,
             data=image[415 : 415 + 474],
         ),
-        timeout_s=2.500,
+        timeout_s=2.5,
     )
 
     # assert that upload() raises SMPUploadError
@@ -386,7 +386,7 @@ async def mock_request(
 @pytest.mark.asyncio
 async def test_upload_file() -> None:
     m = SMPMockTransport()
-    s = SMPClient(m, "address")
+    s = SMPClient(m, "address", 2.5)
 
     s.request = AsyncMock()  # type: ignore
 
@@ -617,7 +617,7 @@ async def mock_request(
 @pytest.mark.asyncio
 async def test_download_file() -> None:
     m = SMPMockTransport()
-    s = SMPClient(m, "address")
+    s = SMPClient(m, "address", 2.5)
 
     s.request = AsyncMock()  # type: ignore