From 97fe0b837c3572b1407ff532b36f52e37c12de27 Mon Sep 17 00:00:00 2001
From: mhidalgo-bdai <144129882+mhidalgo-bdai@users.noreply.github.com>
Date: Mon, 25 Mar 2024 14:55:43 -0300
Subject: [PATCH] Add power service to builtin mocks (#96)

Signed-off-by: Michel Hidalgo <mhidalgo@theaiinstitute.com>
---
 spot_wrapper/testing/mocks/__init__.py |  2 +
 spot_wrapper/testing/mocks/power.py    | 59 ++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 spot_wrapper/testing/mocks/power.py

diff --git a/spot_wrapper/testing/mocks/__init__.py b/spot_wrapper/testing/mocks/__init__.py
index 2344041..f66a719 100644
--- a/spot_wrapper/testing/mocks/__init__.py
+++ b/spot_wrapper/testing/mocks/__init__.py
@@ -94,6 +94,7 @@
 from spot_wrapper.testing.mocks.payload_registration import (
     MockPayloadRegistrationService,
 )
+from spot_wrapper.testing.mocks.power import MockPowerService
 from spot_wrapper.testing.mocks.robot_id import MockRobotIdService
 from spot_wrapper.testing.mocks.robot_state import MockRobotStateService
 from spot_wrapper.testing.mocks.time_sync import MockTimeSyncService
@@ -185,6 +186,7 @@ class MockSpot(
     MockLeaseService,
     MockLicenseService,
     MockPayloadRegistrationService,
+    MockPowerService,
     MockRobotIdService,
     MockRobotStateService,
     MockTimeSyncService,
diff --git a/spot_wrapper/testing/mocks/power.py b/spot_wrapper/testing/mocks/power.py
new file mode 100644
index 0000000..fb8f16f
--- /dev/null
+++ b/spot_wrapper/testing/mocks/power.py
@@ -0,0 +1,59 @@
+# Copyright (c) 2024 Boston Dynamics AI Institute LLC. See LICENSE file for more info.
+
+import grpc
+from bosdyn.api.power_pb2 import (
+    FanPowerCommandFeedbackRequest,
+    FanPowerCommandFeedbackResponse,
+    FanPowerCommandRequest,
+    FanPowerCommandResponse,
+    PowerCommandFeedbackRequest,
+    PowerCommandFeedbackResponse,
+    PowerCommandRequest,
+    PowerCommandResponse,
+    PowerCommandStatus,
+)
+from bosdyn.api.power_service_pb2_grpc import PowerServiceServicer
+from bosdyn.api.robot_state_pb2 import PowerState
+
+from spot_wrapper.testing.mocks.robot_state import MockRobotStateService
+
+
+class MockPowerService(PowerServiceServicer, MockRobotStateService):
+    """A mock Spot power service."""
+
+    def PowerCommand(self, request: PowerCommandRequest, context: grpc.ServicerContext) -> PowerCommandResponse:
+        response = PowerCommandResponse()
+        if request.request == PowerCommandRequest.Request.REQUEST_ON_MOTORS:
+            self.robot_state.power_state.motor_power_state = PowerState.MotorPowerState.MOTOR_POWER_STATE_ON
+            response.status = PowerCommandStatus.STATUS_SUCCESS
+            response.power_command_id = 1
+        elif request.request == PowerCommandRequest.Request.REQUEST_OFF_MOTORS:
+            self.robot_state.power_state.motor_power_state = PowerState.MotorPowerState.MOTOR_POWER_STATE_OFF
+            response.status = PowerCommandStatus.STATUS_SUCCESS
+            response.power_command_id = 1
+        else:
+            response.status = PowerCommandStatus.STATUS_INTERNAL_ERROR
+        return response
+
+    def PowerCommandFeedback(
+        self, request: PowerCommandFeedbackRequest, context: grpc.ServicerContext
+    ) -> PowerCommandFeedbackResponse:
+        response = PowerCommandFeedbackResponse()
+        response.status = PowerCommandStatus.STATUS_SUCCESS
+        response.power_command_id = request.power_command_id
+        return response
+
+    def FanPowerCommand(
+        self, request: FanPowerCommandRequest, context: grpc.ServicerContext
+    ) -> FanPowerCommandResponse:
+        response = FanPowerCommandResponse()
+        response.status = FanPowerCommandResponse.Status.STATUS_OK
+        response.command_id = 2
+        return response
+
+    def FanPowerCommandFeedback(
+        self, request: FanPowerCommandFeedbackRequest, context: grpc.ServicerContext
+    ) -> FanPowerCommandFeedbackResponse:
+        response = FanPowerCommandFeedbackResponse()
+        response.status = FanPowerCommandFeedbackResponse.Status.STATUS_COMPLETE
+        return response