diff --git a/src/rivian/rivian.py b/src/rivian/rivian.py index 1905376..cc2442c 100644 --- a/src/rivian/rivian.py +++ b/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 diff --git a/tests/responses.py b/tests/responses.py index d5dc6cc..69c7bd1 100644 --- a/tests/responses.py +++ b/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]: diff --git a/tests/rivian_test.py b/tests/rivian_test.py index 7eb0f76..ee5db91 100644 --- a/tests/rivian_test.py +++ b/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"