Add the ability to skip fetching zones and sensors

pydrawise/client.py

@@ -110,7 +110,7 @@ async def _mutation(self, selector: DSLField) -> None:
     async def get_user(self, fetch_zones: bool = True) -> User:
         """Retrieves the currently authenticated user.
 
-        :param fetch_zones: Not used in this implementation.
+        :param fetch_zones: Whether to include zones in the controller response.
         :rtype: User
         """
         skip = [] if fetch_zones else ["controllers.zones"]
@@ -120,14 +120,24 @@ async def get_user(self, fetch_zones: bool = True) -> User:
         result = await self._query(selector)
         return deserialize(User, result["me"])
 
-    async def get_controllers(self) -> list[Controller]:
+    async def get_controllers(
+        self, fetch_zones: bool = True, fetch_sensors: bool = True
+    ) -> list[Controller]:
         """Retrieves all controllers associated with the currently authenticated user.
 
+        :param fetch_zones: Whether to include zones in the response.
+        :param fetch_sensors: Whether to include sensors in the response.
         :rtype: list[Controller]
         """
+        skip = []
+        if not fetch_zones:
+            skip.append("zones")
+        if not fetch_sensors:
+            skip.append("sensors")
+
         selector = self._schema.Query.me.select(
             self._schema.User.controllers.select(
-                *get_selectors(self._schema, Controller)
+                *get_selectors(self._schema, Controller, skip)
             ),
         )
         result = await self._query(selector)

tests/test_client.py

@@ -350,9 +350,41 @@ async def test_get_controllers(api: Hydrawise, mock_session, controller_json):
     assert query.count("zones {") == 2
     assert controller.last_contact_time == datetime(2023, 1, 1, 0, 0, 0)
     assert controller.last_action == datetime(2023, 1, 1, 0, 0, 0)
+    assert controller.status is not None
     assert controller.status.actual_water_time.value == timedelta(minutes=10)
 
 
+async def test_get_controllers_no_zones(api: Hydrawise, mock_session, controller_json):
+    mock_session.execute.return_value = {"me": {"controllers": [controller_json]}}
+    [controller] = await api.get_controllers(fetch_zones=False)
+    mock_session.execute.assert_awaited_once()
+    [selector] = mock_session.execute.await_args.args
+    query = print_ast(selector)
+    assert query.count("zones {") == 1
+    assert controller.last_contact_time == datetime(2023, 1, 1, 0, 0, 0)
+    assert controller.last_action == datetime(2023, 1, 1, 0, 0, 0)
+    assert controller.status is not None
+    assert controller.status.actual_water_time.value == timedelta(minutes=10)
+    assert len(controller.zones) == 0
+
+
+async def test_get_controllers_no_sensors(
+    api: Hydrawise, mock_session, controller_json
+):
+    del controller_json["sensors"]
+    mock_session.execute.return_value = {"me": {"controllers": [controller_json]}}
+    [controller] = await api.get_controllers(fetch_sensors=False)
+    mock_session.execute.assert_awaited_once()
+    [selector] = mock_session.execute.await_args.args
+    query = print_ast(selector)
+    assert query.count("sensors {") == 0
+    assert controller.last_contact_time == datetime(2023, 1, 1, 0, 0, 0)
+    assert controller.last_action == datetime(2023, 1, 1, 0, 0, 0)
+    assert controller.status is not None
+    assert controller.status.actual_water_time.value == timedelta(minutes=10)
+    assert len(controller.sensors) == 0
+
+
 async def test_get_controller(api: Hydrawise, mock_session, controller_json):
     mock_session.execute.return_value = {"controller": controller_json}
     controller = await api.get_controller(9876)
@@ -365,6 +397,7 @@ async def test_get_controller(api: Hydrawise, mock_session, controller_json):
 
     assert controller.last_contact_time == datetime(2023, 1, 1, 0, 0, 0)
     assert controller.last_action == datetime(2023, 1, 1, 0, 0, 0)
+    assert controller.status is not None
     assert controller.status.actual_water_time.value == timedelta(minutes=10)