diff --git a/pydrawise/client.py b/pydrawise/client.py index 4faa8cb..c6c1958 100644 --- a/pydrawise/client.py +++ b/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) diff --git a/tests/test_client.py b/tests/test_client.py index 8e4c93e..2e08135 100644 --- a/tests/test_client.py +++ b/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)