From f5f0d14117d72d3403c049f724e14e31300eb484 Mon Sep 17 00:00:00 2001 From: Andrew Rankin Date: Thu, 9 Apr 2026 10:03:12 -0400 Subject: [PATCH 1/3] feat: add get_charging_schedules method Add GraphQL query method to fetch vehicle charging schedules, supporting the upcoming charge amperage control feature in the HA integration. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/rivian/rivian.py | 14 ++++++++++++++ tests/responses.py | 26 ++++++++++++++++++++++++++ tests/rivian_test.py | 22 ++++++++++++++++++++++ 3 files changed, 62 insertions(+) diff --git a/src/rivian/rivian.py b/src/rivian/rivian.py index 1905376..05198b6 100644 --- a/src/rivian/rivian.py +++ b/src/rivian/rivian.py @@ -444,6 +444,20 @@ 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 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..7b2c5db 100644 --- a/tests/responses.py +++ b/tests/responses.py @@ -672,6 +672,32 @@ } +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..be94c2b 100644 --- a/tests/rivian_test.py +++ b/tests/rivian_test.py @@ -23,6 +23,7 @@ LIVE_CHARGING_SESSION_RESPONSE, OTP_TOKEN_RESPONSE, USER_INFORMATION_RESPONSE, + VEHICLE_CHARGING_SCHEDULES_RESPONSE, VEHICLE_STATE_RESPONSE, WALLBOXES_RESPONSE, error_response, @@ -194,6 +195,27 @@ 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_graphql_errors(aresponses: ResponsesMockServer) -> None: """Test GraphQL error responses.""" host = "rivian.com" From c697ca70ce6aaa94c90392690993ad4ab50c9e5c Mon Sep 17 00:00:00 2001 From: Andrew Rankin Date: Thu, 9 Apr 2026 10:21:41 -0400 Subject: [PATCH 2/3] feat: add set_charging_schedules method Co-Authored-By: Claude Opus 4.6 (1M context) --- src/rivian/rivian.py | 20 ++++++++++++++++++++ tests/responses.py | 9 +++++++++ tests/rivian_test.py | 40 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 69 insertions(+) diff --git a/src/rivian/rivian.py b/src/rivian/rivian.py index 05198b6..5856434 100644 --- a/src/rivian/rivian.py +++ b/src/rivian/rivian.py @@ -458,6 +458,26 @@ async def get_charging_schedules(self, vehicle_id: str) -> ClientResponse: } 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 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 7b2c5db..69c7bd1 100644 --- a/tests/responses.py +++ b/tests/responses.py @@ -672,6 +672,15 @@ } +SET_CHARGING_SCHEDULES_RESPONSE = { + "data": { + "setChargingSchedules": { + "success": True + } + } +} + + VEHICLE_CHARGING_SCHEDULES_RESPONSE = { "data": { "getVehicle": { diff --git a/tests/rivian_test.py b/tests/rivian_test.py index be94c2b..ee5db91 100644 --- a/tests/rivian_test.py +++ b/tests/rivian_test.py @@ -22,6 +22,7 @@ CSRF_TOKEN_RESPONSE, LIVE_CHARGING_SESSION_RESPONSE, OTP_TOKEN_RESPONSE, + SET_CHARGING_SCHEDULES_RESPONSE, USER_INFORMATION_RESPONSE, VEHICLE_CHARGING_SCHEDULES_RESPONSE, VEHICLE_STATE_RESPONSE, @@ -216,6 +217,45 @@ async def test_get_charging_schedules(aresponses: ResponsesMockServer) -> None: 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" From 4784ec58aa626d56c56706f236567b3dc6fb6427 Mon Sep 17 00:00:00 2001 From: Andrew Rankin Date: Thu, 9 Apr 2026 10:28:29 -0400 Subject: [PATCH 3/3] style: add __typename to setChargingSchedules mutation response Matches the convention used by all other mutations in the codebase. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/rivian/rivian.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rivian/rivian.py b/src/rivian/rivian.py index 5856434..cc2442c 100644 --- a/src/rivian/rivian.py +++ b/src/rivian/rivian.py @@ -470,7 +470,7 @@ async def set_charging_schedules( } graphql_json = { "operationName": "setChargingSchedules", - "query": "mutation setChargingSchedules($vehicleId: String!, $chargingSchedules: [InputChargingSchedule!]!) {\n setChargingSchedules(vehicleId: $vehicleId, chargingSchedules: $chargingSchedules) {\n success\n }\n}", + "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,