RSDK-11344 - implement restart_module (viamrobotics#951)

Author: stuqdog

src/viam/robot/client.py

@@ -35,6 +35,7 @@
     Operation,
     ResourceNamesRequest,
     ResourceNamesResponse,
+    RestartModuleRequest,
     RobotServiceStub,
     ShutdownRequest,
     StopAllRequest,
@@ -907,3 +908,31 @@ async def get_machine_status(self) -> GetMachineStatusResponse:
 
         request = GetMachineStatusRequest()
         return await self._client.GetMachineStatus(request)
+
+
+    ##################
+    # Restart Module #
+    ##################
+
+    async def restart_module(self, id: Optional[str] = None, name: Optional[str] = None):
+        """
+        Restarts a module running on the machine with the given id or name.
+
+        ::
+
+            await machine.restart_module(id="namespace:module:model", name="my_model")
+
+        Args:
+            id (str): The id matching the module_id field of the registry module in your part configuration.
+            name (str): The name matching the name field of the local/registry module in your part configuration.
+
+        Raises:
+            GRPCError: If a module can't be found matching the provided ID or name.
+
+        For more information, see `Machine Management API <https://docs.viam.com/appendix/apis/robot/>`_.
+        """
+
+        id = id if id else ""
+        name = name if name else ""
+        request = RestartModuleRequest(module_id=id, module_name=name)
+        await self._client.RestartModule(request)

tests/test_board.py

@@ -361,8 +361,9 @@ async def test_set_power_mode(self, board: MockBoard, service: BoardRPCService):
             client = BoardClient(name=board.name, channel=channel)
             pm_mode = PowerMode.POWER_MODE_OFFLINE_DEEP
             pm_timedelta = timedelta(minutes=1)
-            await client.set_power_mode(mode=pm_mode, duration=pm_timedelta, timeout=1.1)
-            assert board.timeout == loose_approx(1.1)
+            await client.set_power_mode(mode=pm_mode, duration=pm_timedelta, timeout=9.83)
+            print("timeout is", board.timeout)
+            assert board.timeout == loose_approx(9.83)
             assert board.power_mode == pm_mode
             pm_duration = Duration()
             pm_duration.FromTimedelta(pm_timedelta)

tests/test_robot.py

@@ -32,6 +32,8 @@
     GetVersionRequest,
     GetVersionResponse,
     Operation,
+    RestartModuleRequest,
+    RestartModuleResponse,
     ResourceNamesRequest,
     ResourceNamesResponse,
     ResourceStatus,
@@ -92,6 +94,9 @@
     machine_part_id="the-machine-part-id",
 )
 
+MODULE_ID = "id"
+MODULE_NAME = "name"
+
 GET_VERVSION_RESPONSE = GetVersionResponse(
     platform="rdk",
     version="0.2.0",
@@ -176,6 +181,13 @@ async def Shutdown(stream: Stream[ShutdownRequest, ShutdownResponse]) -> None:
         response = ShutdownResponse()
         await stream.send_message(response)
 
+    async def RestartModule(stream: Stream[RestartModuleRequest, RestartModuleResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        assert request.module_name == MODULE_NAME
+        assert request.module_id == MODULE_ID
+        await stream.send_message(RestartModuleResponse())
+
     manager = ResourceManager(resources)
     service = RobotService(manager)
     service.FrameSystemConfig = Config
@@ -549,3 +561,14 @@ async def shutdown_client_mock(self):
                 shutdown_mock.assert_called_once()
 
                 await client.close()
+
+    async def test_restart_module(self, service: RobotService):
+        async with ChannelFor([service]) as channel:
+
+            client = await RobotClient.with_channel(channel, RobotClient.Options())
+
+            with mock.patch("viam.robot.client.RobotClient.restart_module") as restart_module_mock:
+                await client.restart_module(id=MODULE_ID, name=MODULE_NAME)
+                restart_module_mock.assert_called_once()
+
+                await client.close()