From d9e500daadef3449b3c7a2ec18f4ebe71b35ad08 Mon Sep 17 00:00:00 2001
From: Michel Hidalgo <mhidalgo@theaiinstitute.com>
Date: Tue, 25 Jun 2024 18:44:20 -0300
Subject: [PATCH] Add basic Spot CAM mocks

Signed-off-by: Michel Hidalgo <mhidalgo@theaiinstitute.com>
---
 spot_wrapper/testing/mocks/__init__.py        |  8 ++++----
 spot_wrapper/testing/mocks/cam.py             | 20 +++++++++++++++++++
 .../{payload_registration.py => payload.py}   | 18 ++++++++++-------
 3 files changed, 35 insertions(+), 11 deletions(-)
 create mode 100644 spot_wrapper/testing/mocks/cam.py
 rename spot_wrapper/testing/mocks/{payload_registration.py => payload.py} (80%)

diff --git a/spot_wrapper/testing/mocks/__init__.py b/spot_wrapper/testing/mocks/__init__.py
index f66a7195..eca51007 100644
--- a/spot_wrapper/testing/mocks/__init__.py
+++ b/spot_wrapper/testing/mocks/__init__.py
@@ -86,14 +86,13 @@
 from spot_wrapper.testing.grpc import AutoServicer, collect_method_handlers
 from spot_wrapper.testing.helpers import enforce_matching_headers
 from spot_wrapper.testing.mocks.auth import MockAuthService
+from spot_wrapper.testing.mocks.cam import MockCAMService
 from spot_wrapper.testing.mocks.directory import MockDirectoryService
 from spot_wrapper.testing.mocks.estop import MockEStopService
 from spot_wrapper.testing.mocks.keepalive import MockKeepaliveService
 from spot_wrapper.testing.mocks.lease import MockLeaseService
 from spot_wrapper.testing.mocks.license import MockLicenseService
-from spot_wrapper.testing.mocks.payload_registration import (
-    MockPayloadRegistrationService,
-)
+from spot_wrapper.testing.mocks.payload import MockPayloadService
 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
@@ -180,12 +179,13 @@ def __init__(self, *args: typing.Any, **kwargs: typing.Any) -> None:
 class MockSpot(
     BaseMockSpot,
     MockAuthService,
+    MockCAMService,
     MockDirectoryService,
     MockEStopService,
     MockKeepaliveService,
     MockLeaseService,
     MockLicenseService,
-    MockPayloadRegistrationService,
+    MockPayloadService,
     MockPowerService,
     MockRobotIdService,
     MockRobotStateService,
diff --git a/spot_wrapper/testing/mocks/cam.py b/spot_wrapper/testing/mocks/cam.py
new file mode 100644
index 00000000..ee1f439c
--- /dev/null
+++ b/spot_wrapper/testing/mocks/cam.py
@@ -0,0 +1,20 @@
+# Copyright (c) 2024 Boston Dynamics AI Institute LLC. See LICENSE file for more info.
+
+import typing
+
+import grpc
+from bosdyn.api.spot_cam.ptz_pb2 import ListPtzRequest, ListPtzResponse, PtzDescription
+from bosdyn.api.spot_cam.service_pb2_grpc import PtzServiceServicer
+
+
+class MockCAMService(PtzServiceServicer):
+    """Mock Spot CAM services."""
+
+    def __init__(self, **kwargs: typing.Any) -> None:
+        super().__init__(**kwargs)
+        self.ptzs: typing.List[PtzDescription] = []
+
+    def ListPtz(self, request: ListPtzRequest, context: grpc.ServicerContext) -> ListPtzResponse:
+        response = ListPtzResponse()
+        response.ptzs.extend(self.ptzs)
+        return response
diff --git a/spot_wrapper/testing/mocks/payload_registration.py b/spot_wrapper/testing/mocks/payload.py
similarity index 80%
rename from spot_wrapper/testing/mocks/payload_registration.py
rename to spot_wrapper/testing/mocks/payload.py
index 54aa1419..b4d1b2a1 100644
--- a/spot_wrapper/testing/mocks/payload_registration.py
+++ b/spot_wrapper/testing/mocks/payload.py
@@ -3,7 +3,7 @@
 import typing
 
 import grpc
-from bosdyn.api.payload_pb2 import Payload
+from bosdyn.api.payload_pb2 import ListPayloadsRequest, ListPayloadsResponse, Payload
 from bosdyn.api.payload_registration_pb2 import (
     GetPayloadAuthTokenRequest,
     GetPayloadAuthTokenResponse,
@@ -17,22 +17,26 @@
 from bosdyn.api.payload_registration_service_pb2_grpc import (
     PayloadRegistrationServiceServicer,
 )
+from bosdyn.api.payload_service_pb2_grpc import PayloadServiceServicer
 
 
-class MockPayloadRegistrationService(PayloadRegistrationServiceServicer):
+class MockPayloadService(PayloadRegistrationServiceServicer, PayloadServiceServicer):
     """
-    A mock Spot payload registration service.
+    A mock Spot payload registration and listing service.
 
     It bookkeeps all payloads but enforces nothing.
     """
 
     def __init__(self, **kwargs: typing.Any) -> None:
         super().__init__(**kwargs)
-        self._payloads: typing.Dict[str, Payload] = {}
+        self.payloads: typing.Dict[str, Payload] = {
+            "spotcam": Payload(GUID="spotcam", name="Spot CAM", is_enabled=True)
+        }
 
-    @property
-    def payloads(self) -> typing.Iterable[Payload]:
-        return list(self._payloads.values())
+    def ListPayloads(self, request: ListPayloadsRequest, context: grpc.ServicerContext) -> ListPayloadsResponse:
+        response = ListPayloadsResponse()
+        response.payloads.extend(self.payloads.values())
+        return response
 
     def RegisterPayload(
         self, request: RegisterPayloadRequest, context: grpc.ServicerContext