feat: add charging schedule API methods

src/rivian/rivian.py

@@ -444,6 +444,40 @@ async def get_vehicle_state(
 
         return await self.__graphql_query(headers, url, graphql_json)
 
+    async def get_charging_schedules(self, vehicle_id: str) -> ClientResponse:
+        """Get charging schedules for a vehicle."""
+        url = GRAPHQL_GATEWAY
+        headers = BASE_HEADERS | {
+            "A-Sess": self._app_session_token,
+            "U-Sess": self._user_session_token,
+        }
+        graphql_json = {
+            "operationName": "getVehicleChargingSchedules",
+            "query": "query getVehicleChargingSchedules($vehicleId: String!) {\n  getVehicle(id: $vehicleId) {\n    chargingSchedules {\n      weekDays\n      startTime\n      duration\n      location {\n        latitude\n        longitude\n      }\n      amperage\n      enabled\n    }\n  }\n}",
+            "variables": {"vehicleId": vehicle_id},
+        }
+        return await self.__graphql_query(headers, url, graphql_json)
+
+    async def set_charging_schedules(
+        self, vehicle_id: str, schedules: list[dict[str, Any]]
+    ) -> ClientResponse:
+        """Set charging schedules for a vehicle."""
+        url = GRAPHQL_GATEWAY
+        headers = BASE_HEADERS | {
+            "Csrf-Token": self._csrf_token,
+            "A-Sess": self._app_session_token,
+            "U-Sess": self._user_session_token,
+        }
+        graphql_json = {
+            "operationName": "setChargingSchedules",
+            "query": "mutation setChargingSchedules($vehicleId: String!, $chargingSchedules: [InputChargingSchedule!]!) {\n  setChargingSchedules(vehicleId: $vehicleId, chargingSchedules: $chargingSchedules) {\n    __typename\n    success\n  }\n}",
+            "variables": {
+                "vehicleId": vehicle_id,
+                "chargingSchedules": schedules,
+            },
+        }
+        return await self.__graphql_query(headers, url, graphql_json)
+
     async def get_vehicle_ota_update_details(self, vehicle_id: str) -> ClientResponse:
         """Get vehicle OTA update details."""
         url = GRAPHQL_GATEWAY

tests/responses.py

@@ -672,6 +672,41 @@
 }
 
 
+SET_CHARGING_SCHEDULES_RESPONSE = {
+    "data": {
+        "setChargingSchedules": {
+            "success": True
+        }
+    }
+}
+
+
+VEHICLE_CHARGING_SCHEDULES_RESPONSE = {
+    "data": {
+        "getVehicle": {
+            "chargingSchedules": [
+                {
+                    "weekDays": [
+                        "Monday",
+                        "Tuesday",
+                        "Wednesday",
+                        "Thursday",
+                        "Friday",
+                        "Saturday",
+                        "Sunday",
+                    ],
+                    "startTime": 0,
+                    "duration": 1440,
+                    "location": {"latitude": 37.7749, "longitude": -122.4194},
+                    "amperage": 32,
+                    "enabled": True,
+                }
+            ]
+        }
+    }
+}
+
+
 def error_response(
     code: str | None = None, reason: str | None = None
 ) -> dict[str, Any]:

tests/rivian_test.py

@@ -22,7 +22,9 @@
     CSRF_TOKEN_RESPONSE,
     LIVE_CHARGING_SESSION_RESPONSE,
     OTP_TOKEN_RESPONSE,
+    SET_CHARGING_SCHEDULES_RESPONSE,
     USER_INFORMATION_RESPONSE,
+    VEHICLE_CHARGING_SCHEDULES_RESPONSE,
     VEHICLE_STATE_RESPONSE,
     WALLBOXES_RESPONSE,
     error_response,
@@ -194,6 +196,66 @@ async def test_get_live_charging_session(aresponses: ResponsesMockServer) -> Non
         await rivian.close()
 
 
+async def test_get_charging_schedules(aresponses: ResponsesMockServer) -> None:
+    """Test getting vehicle charging schedules."""
+    aresponses.add(
+        "rivian.com",
+        "/api/gql/gateway/graphql",
+        "POST",
+        response=VEHICLE_CHARGING_SCHEDULES_RESPONSE,
+    )
+    async with aiohttp.ClientSession():
+        rivian = Rivian(app_session_token="token", user_session_token="token")
+        response = await rivian.get_charging_schedules("vehicle_id")
+        response_json = await response.json()
+        assert response.status == 200
+        schedules = response_json["data"]["getVehicle"]["chargingSchedules"]
+        assert len(schedules) == 1
+        assert schedules[0]["amperage"] == 32
+        assert schedules[0]["enabled"] is True
+        assert len(schedules[0]["weekDays"]) == 7
+        await rivian.close()
+
+
+async def test_set_charging_schedules(aresponses: ResponsesMockServer) -> None:
+    """Test setting vehicle charging schedules."""
+    aresponses.add(
+        "rivian.com",
+        "/api/gql/gateway/graphql",
+        "POST",
+        response=SET_CHARGING_SCHEDULES_RESPONSE,
+    )
+    async with aiohttp.ClientSession():
+        rivian = Rivian(
+            csrf_token="csrf",
+            app_session_token="token",
+            user_session_token="token",
+        )
+        schedules = [
+            {
+                "weekDays": [
+                    "Monday",
+                    "Tuesday",
+                    "Wednesday",
+                    "Thursday",
+                    "Friday",
+                    "Saturday",
+                    "Sunday",
+                ],
+                "startTime": 0,
+                "duration": 1440,
+                "location": {"latitude": 37.7749, "longitude": -122.4194},
+                "amperage": 16,
+                "enabled": True,
+            }
+        ]
+        response = await rivian.set_charging_schedules("vehicle_id", schedules)
+        response_json = await response.json()
+        assert response.status == 200
+        assert response_json["data"]["setChargingSchedules"]["success"] is True
+        await rivian.close()
+
+
 async def test_graphql_errors(aresponses: ResponsesMockServer) -> None:
     """Test GraphQL error responses."""
     host = "rivian.com"