Expose builtin mocks state (#97)

Signed-off-by: Michel Hidalgo <[email protected]>
bdaiinstitute · Mar 25, 2024 · e238ad8 · e238ad8
1 parent 97fe0b8
commit e238ad8
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 7 deletions.
diff --git a/spot_wrapper/testing/mocks/lease.py b/spot_wrapper/testing/mocks/lease.py
@@ -55,6 +55,10 @@ def __init__(
             self._leasable_resources[resource.resource] = leasable_resource
         self._latest_lease: typing.Optional[Lease] = None
 
+    @property
+    def leasable_resources(self) -> typing.Iterable[LeaseResource]:
+        return list(self._leasable_resources.values())
+
     def AcquireLease(self, request: AcquireLeaseRequest, context: grpc.ServicerContext) -> AcquireLeaseResponse:
         response = AcquireLeaseResponse()
         if request.resource not in self._leasable_resources:

diff --git a/spot_wrapper/testing/mocks/payload_registration.py b/spot_wrapper/testing/mocks/payload_registration.py
@@ -30,6 +30,10 @@ def __init__(self, **kwargs: typing.Any) -> None:
         super().__init__(**kwargs)
         self._payloads: typing.Dict[str, Payload] = {}
 
+    @property
+    def payloads(self) -> typing.Iterable[Payload]:
+        return list(self._payloads.values())
+
     def RegisterPayload(
         self, request: RegisterPayloadRequest, context: grpc.ServicerContext
     ) -> RegisterPayloadResponse:

diff --git a/spot_wrapper/testing/mocks/robot_state.py b/spot_wrapper/testing/mocks/robot_state.py
@@ -28,8 +28,8 @@ class MockRobotStateService(RobotStateServiceServicer):
 
     def __init__(self, **kwargs: typing.Any) -> None:
         super().__init__(**kwargs)
-        self.robot_state = RobotState()
-        transforms_snapshot = self.robot_state.kinematic_state.transforms_snapshot
+        self._robot_state = RobotState()
+        transforms_snapshot = self._robot_state.kinematic_state.transforms_snapshot
         world_to_odom_edge = transforms_snapshot.child_to_parent_edge_map["odom"]
         world_to_odom_edge.parent_tform_child.rotation.w = 1.0
         odom_to_body_edge = transforms_snapshot.child_to_parent_edge_map["body"]
@@ -38,24 +38,36 @@ def __init__(self, **kwargs: typing.Any) -> None:
         body_to_vision_edge = transforms_snapshot.child_to_parent_edge_map["vision"]
         body_to_vision_edge.parent_frame_name = "body"
         body_to_vision_edge.parent_tform_child.rotation.w = 1.0
-        self.robot_metrics = RobotMetrics()
-        self.hardware_configuration = HardwareConfiguration()
+        self._robot_metrics = RobotMetrics()
+        self._hardware_configuration = HardwareConfiguration()
+
+    @property
+    def robot_state(self) -> RobotState:
+        return self._robot_state
+
+    @property
+    def robot_metrics(self) -> RobotMetrics:
+        return self._robot_metrics
+
+    @property
+    def hardware_configuration(self) -> HardwareConfiguration:
+        return self._hardware_configuration
 
     def GetRobotState(self, request: RobotStateRequest, context: grpc.ServicerContext) -> RobotStateResponse:
         response = RobotStateResponse()
-        response.robot_state.CopyFrom(self.robot_state)
+        response.robot_state.CopyFrom(self._robot_state)
         return response
 
     def GetRobotMetrics(self, request: RobotMetricsRequest, context: grpc.ServicerContext) -> RobotMetricsResponse:
         response = RobotMetricsResponse()
-        response.robot_metrics.CopyFrom(self.robot_metrics)
+        response.robot_metrics.CopyFrom(self._robot_metrics)
         return response
 
     def GetRobotHardwareConfiguration(
         self, request: RobotHardwareConfigurationRequest, context: grpc.ServicerContext
     ) -> RobotHardwareConfigurationResponse:
         response = RobotHardwareConfigurationResponse()
-        response.hardware_configuration.CopyFrom(self.hardware_configuration)
+        response.hardware_configuration.CopyFrom(self._hardware_configuration)
         return response
 
     def GetRobotLinkModel(