RSDK-7698: Add Shutdown command to Python SDK (#633)

Author: danielbotros

src/viam/robot/client.py

@@ -4,6 +4,7 @@
 from threading import RLock
 from typing import Any, Dict, List, Optional, Union
 
+from grpclib import GRPCError, Status
 from grpclib.client import Channel
 from typing_extensions import Self
 
@@ -33,6 +34,7 @@
     ResourceNamesRequest,
     ResourceNamesResponse,
     RobotServiceStub,
+    ShutdownRequest,
     StopAllRequest,
     StopExtraParameters,
     TransformPoseRequest,
@@ -787,3 +789,33 @@ async def get_cloud_metadata(self) -> GetCloudMetadataResponse:
 
         request = GetCloudMetadataRequest()
         return await self._client.GetCloudMetadata(request)
+
+    ############
+    # Shutdown #
+    ############
+
+    async def shutdown(self):
+        """
+        Shutdown shuts down the robot.
+
+        Raises:
+            GRPCError: Raised with DeadlineExceeded status if shutdown request times out, or if
+              robot server shuts down before having a chance to send a response. Raised with
+              status Unavailable if server is unavailable, or if robot server is in the process of
+              shutting down when response is ready.
+        """
+        request = ShutdownRequest()
+        try:
+            await self._client.Shutdown(request)
+            LOGGER.info("robot shutdown successful")
+        except GRPCError as e:
+            if e.status == Status.INTERNAL or e.status == Status.UNKNOWN:
+                LOGGER.info("robot shutdown successful")
+            elif e.status == Status.UNAVAILABLE:
+                LOGGER.warn("server unavailable, likely due to successful robot shutdown")
+                raise e
+            elif e.status == Status.DEADLINE_EXCEEDED:
+                LOGGER.warn("request timeout, robot shutdown may still be successful")
+                raise e
+            else:
+                raise e

tests/mocks/robot.py

@@ -31,6 +31,8 @@
     RestartModuleResponse,
     SendSessionHeartbeatRequest,
     SendSessionHeartbeatResponse,
+    ShutdownRequest,
+    ShutdownResponse,
     StartSessionRequest,
     StartSessionResponse,
     StopAllRequest,
@@ -115,3 +117,6 @@ async def Log(self, stream: Stream[LogRequest, LogResponse]) -> None:
 
     async def GetCloudMetadata(self, stream: Stream[GetCloudMetadataRequest, GetCloudMetadataResponse]) -> None:
         raise MethodNotImplementedError("GetCloudMetadata").grpc_error
+
+    async def Shutdown(self, stream: Stream[ShutdownRequest, ShutdownResponse]) -> None:
+        raise MethodNotImplementedError("Shutdown").grpc_error

tests/test_robot.py

@@ -1,5 +1,6 @@
 import asyncio
 from typing import Any, Dict, List, Optional, Tuple
+from unittest import mock
 
 import pytest
 from google.protobuf.struct_pb2 import Struct, Value
@@ -39,6 +40,8 @@
     ResourceNamesRequest,
     ResourceNamesResponse,
     RobotServiceStub,
+    ShutdownRequest,
+    ShutdownResponse,
     Status,
     StopAllRequest,
     StopExtraParameters,
@@ -209,13 +212,20 @@ async def GetCloudMetadata(stream: Stream[GetCloudMetadataRequest, GetCloudMetad
         assert request is not None
         await stream.send_message(GET_CLOUD_METADATA_RESPONSE)
 
+    async def Shutdown(stream: Stream[ShutdownRequest, ShutdownResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        response = ShutdownResponse()
+        await stream.send_message(response)
+
     manager = ResourceManager(resources)
     service = RobotService(manager)
     service.FrameSystemConfig = Config
     service.TransformPose = TransformPose
     service.DiscoverComponents = DiscoverComponents
     service.GetOperations = GetOperations
     service.GetCloudMetadata = GetCloudMetadata
+    service.Shutdown = Shutdown
 
     return service
 
@@ -599,3 +609,21 @@ async def get_geometries(
 
             await client.close()
             Registry._SUBTYPES[Arm.SUBTYPE].create_rpc_client = old_create_client
+
+    @pytest.mark.asyncio
+    async def test_shutdown(self, service: RobotService):
+        async with ChannelFor([service]) as channel:
+
+            async def shutdown_client_mock(self):
+                return await self._client.Shutdown(ShutdownRequest())
+
+            client = await RobotClient.with_channel(channel, RobotClient.Options())
+
+            with mock.patch("viam.robot.client.RobotClient.shutdown") as shutdown_mock:
+                shutdown_mock.return_value = await shutdown_client_mock(client)
+                shutdown_response = await client.shutdown()
+
+                assert shutdown_response == ShutdownResponse()
+                shutdown_mock.assert_called_once()
+
+                await client.close()