Skip to content

Commit 4d6853f

Browse files
authored
RSDK-11344 - implement restart_module (#951)
1 parent f23f87f commit 4d6853f

File tree

3 files changed

+55
-2
lines changed

3 files changed

+55
-2
lines changed

src/viam/robot/client.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
Operation,
3636
ResourceNamesRequest,
3737
ResourceNamesResponse,
38+
RestartModuleRequest,
3839
RobotServiceStub,
3940
ShutdownRequest,
4041
StopAllRequest,
@@ -907,3 +908,31 @@ async def get_machine_status(self) -> GetMachineStatusResponse:
907908

908909
request = GetMachineStatusRequest()
909910
return await self._client.GetMachineStatus(request)
911+
912+
913+
##################
914+
# Restart Module #
915+
##################
916+
917+
async def restart_module(self, id: Optional[str] = None, name: Optional[str] = None):
918+
"""
919+
Restarts a module running on the machine with the given id or name.
920+
921+
::
922+
923+
await machine.restart_module(id="namespace:module:model", name="my_model")
924+
925+
Args:
926+
id (str): The id matching the module_id field of the registry module in your part configuration.
927+
name (str): The name matching the name field of the local/registry module in your part configuration.
928+
929+
Raises:
930+
GRPCError: If a module can't be found matching the provided ID or name.
931+
932+
For more information, see `Machine Management API <https://docs.viam.com/appendix/apis/robot/>`_.
933+
"""
934+
935+
id = id if id else ""
936+
name = name if name else ""
937+
request = RestartModuleRequest(module_id=id, module_name=name)
938+
await self._client.RestartModule(request)

tests/test_board.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -361,8 +361,9 @@ async def test_set_power_mode(self, board: MockBoard, service: BoardRPCService):
361361
client = BoardClient(name=board.name, channel=channel)
362362
pm_mode = PowerMode.POWER_MODE_OFFLINE_DEEP
363363
pm_timedelta = timedelta(minutes=1)
364-
await client.set_power_mode(mode=pm_mode, duration=pm_timedelta, timeout=1.1)
365-
assert board.timeout == loose_approx(1.1)
364+
await client.set_power_mode(mode=pm_mode, duration=pm_timedelta, timeout=9.83)
365+
print("timeout is", board.timeout)
366+
assert board.timeout == loose_approx(9.83)
366367
assert board.power_mode == pm_mode
367368
pm_duration = Duration()
368369
pm_duration.FromTimedelta(pm_timedelta)

tests/test_robot.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
GetVersionRequest,
3333
GetVersionResponse,
3434
Operation,
35+
RestartModuleRequest,
36+
RestartModuleResponse,
3537
ResourceNamesRequest,
3638
ResourceNamesResponse,
3739
ResourceStatus,
@@ -92,6 +94,9 @@
9294
machine_part_id="the-machine-part-id",
9395
)
9496

97+
MODULE_ID = "id"
98+
MODULE_NAME = "name"
99+
95100
GET_VERVSION_RESPONSE = GetVersionResponse(
96101
platform="rdk",
97102
version="0.2.0",
@@ -176,6 +181,13 @@ async def Shutdown(stream: Stream[ShutdownRequest, ShutdownResponse]) -> None:
176181
response = ShutdownResponse()
177182
await stream.send_message(response)
178183

184+
async def RestartModule(stream: Stream[RestartModuleRequest, RestartModuleResponse]) -> None:
185+
request = await stream.recv_message()
186+
assert request is not None
187+
assert request.module_name == MODULE_NAME
188+
assert request.module_id == MODULE_ID
189+
await stream.send_message(RestartModuleResponse())
190+
179191
manager = ResourceManager(resources)
180192
service = RobotService(manager)
181193
service.FrameSystemConfig = Config
@@ -549,3 +561,14 @@ async def shutdown_client_mock(self):
549561
shutdown_mock.assert_called_once()
550562

551563
await client.close()
564+
565+
async def test_restart_module(self, service: RobotService):
566+
async with ChannelFor([service]) as channel:
567+
568+
client = await RobotClient.with_channel(channel, RobotClient.Options())
569+
570+
with mock.patch("viam.robot.client.RobotClient.restart_module") as restart_module_mock:
571+
await client.restart_module(id=MODULE_ID, name=MODULE_NAME)
572+
restart_module_mock.assert_called_once()
573+
574+
await client.close()

0 commit comments

Comments
 (0)