finish cover api and clean up tests

Author: dmulcahey

tests/test_cover.py

@@ -94,22 +94,21 @@
 
 
 @pytest.mark.parametrize(
-    "gateway_type, entity_type",
+    "gateway_type",
     [
-        ("zha_gateway", Platform.COVER),
-        ("ws_gateway", Platform.COVER),
+        "zha_gateway",
+        "ws_gateway",
     ],
 )
 @pytest.mark.looptime
 async def test_cover_non_tilt_initial_state(  # pylint: disable=unused-argument
     zha_gateways: CombinedGateways,
-    zigpy_cover_device,
     gateway_type: str,
-    entity_type: type,
 ) -> None:
     """Test ZHA cover platform."""
 
     zha_gateway = getattr(zha_gateways, gateway_type)
+    zigpy_cover_device = create_mock_zigpy_device(zha_gateway, ZIGPY_COVER_DEVICE)
     # load up cover domain
     zigpy_cover_device = create_mock_zigpy_device(zha_gateway, ZIGPY_COVER_DEVICE)
     cluster = zigpy_cover_device.endpoints[1].window_covering
@@ -164,17 +163,16 @@ async def test_cover_non_tilt_initial_state(  # pylint: disable=unused-argument
 
 
 @pytest.mark.parametrize(
-    "gateway_type, entity_type",
+    "gateway_type",
     [
-        ("zha_gateway", Platform.COVER),
-        ("ws_gateway", Platform.COVER),
+        "zha_gateway",
+        "ws_gateway",
     ],
 )
 @pytest.mark.looptime
 async def test_cover(
     zha_gateways: CombinedGateways,
     gateway_type: str,
-    entity_type: type,
 ) -> None:
     """Test zha cover platform."""
 
@@ -410,21 +408,21 @@ async def test_cover(
 
 
 @pytest.mark.parametrize(
-    "gateway_type, entity_type",
+    "gateway_type",
     [
-        ("zha_gateway", Platform.COVER),
-        ("ws_gateway", Platform.COVER),
+        "zha_gateway",
+        "ws_gateway",
     ],
 )
 @pytest.mark.looptime
 async def test_cover_failures(
     zha_gateways: CombinedGateways,
     gateway_type: str,
-    entity_type: type,
 ) -> None:
     """Test ZHA cover platform failure cases."""
 
     zha_gateway = getattr(zha_gateways, gateway_type)
+    zigpy_cover_device = create_mock_zigpy_device(zha_gateway, ZIGPY_COVER_DEVICE)
     # load up cover domain
     zigpy_cover_device = create_mock_zigpy_device(zha_gateway, ZIGPY_COVER_DEVICE)
     cluster = zigpy_cover_device.endpoints[1].window_covering
@@ -621,17 +619,16 @@ async def test_cover_failures(
 
 
 @pytest.mark.parametrize(
-    "gateway_type, entity_type",
+    "gateway_type",
     [
-        ("zha_gateway", Platform.COVER),
-        ("ws_gateway", Platform.COVER),
+        "zha_gateway",
+        "ws_gateway",
     ],
 )
 @pytest.mark.looptime
 async def test_shade(
     zha_gateways: CombinedGateways,
     gateway_type: str,
-    entity_type: type,
 ) -> None:
     """Test zha cover platform for shade device type."""
 
@@ -813,17 +810,16 @@ async def test_shade(
 
 
 @pytest.mark.parametrize(
-    "gateway_type, entity_type",
+    "gateway_type",
     [
-        ("zha_gateway", Platform.COVER),
-        ("ws_gateway", Platform.COVER),
+        "zha_gateway",
+        "ws_gateway",
     ],
 )
 @pytest.mark.looptime
 async def test_keen_vent(
     zha_gateways: CombinedGateways,
     gateway_type: str,
-    entity_type: type,
 ) -> None:
     """Test keen vent."""
 
@@ -893,21 +889,21 @@ async def test_keen_vent(
 
 
 @pytest.mark.parametrize(
-    "gateway_type, entity_type",
+    "gateway_type",
     [
-        ("zha_gateway", Platform.COVER),
-        ("ws_gateway", Platform.COVER),
+        "zha_gateway",
+        "ws_gateway",
     ],
 )
 @pytest.mark.looptime
 async def test_cover_remote(
     zha_gateways: CombinedGateways,
     gateway_type: str,
-    entity_type: type,
 ) -> None:
     """Test ZHA cover remote."""
 
     zha_gateway = getattr(zha_gateways, gateway_type)
+    zigpy_cover_remote = create_mock_zigpy_device(zha_gateway, ZIGPY_COVER_REMOTE)
     # load up cover domain
     zigpy_cover_remote = create_mock_zigpy_device(zha_gateway, ZIGPY_COVER_REMOTE)
     zha_device = await join_zigpy_device(zha_gateway, zigpy_cover_remote)
@@ -946,12 +942,21 @@ async def test_cover_remote(
     assert device.emit_zha_event.call_args[0][0][ATTR_COMMAND] == "down_close"
 
 
-# TODO parametrize this test and add service to restore state attributes
+@pytest.mark.parametrize(
+    "gateway_type",
+    [
+        "zha_gateway",
+        "ws_gateway",
+    ],
+)
 @pytest.mark.looptime
 async def test_cover_state_restoration(
-    zha_gateway: Gateway,
+    zha_gateways: CombinedGateways,
+    gateway_type: str,
 ) -> None:
     """Test the cover state restoration."""
+
+    zha_gateway = getattr(zha_gateways, gateway_type)
     zigpy_cover_device = create_mock_zigpy_device(zha_gateway, ZIGPY_COVER_DEVICE)
     zha_device = await join_zigpy_device(zha_gateway, zigpy_cover_device)
     entity = get_entity(zha_device, platform=Platform.COVER)
@@ -966,6 +971,11 @@ async def test_cover_state_restoration(
         target_tilt_position=34,
     )
 
+    # ws impl needs a round trip to get the state back to the client
+    # maybe we make this optimistic, set the state manually on the client
+    # and avoid the round trip refresh call?
+    await zha_gateway.async_block_till_done()
+
     assert entity.state["state"] == STATE_CLOSED
     assert entity.state["target_lift_position"] == 12
     assert entity.state["target_tilt_position"] == 34

zha/application/platforms/cover/__init__.py

@@ -191,6 +191,7 @@ def restore_external_state_attributes(
         ],  # FIXME: why must these be expanded?
         target_lift_position: int | None,
         target_tilt_position: int | None,
+        **kwargs: Any,
     ):
         """Restore external state attributes."""
         self._state = state
@@ -624,6 +625,7 @@ def __init__(
     ) -> None:
         """Initialize the ZHA fan entity."""
         super().__init__(entity_info, device)
+        self._tasks: list[asyncio.Task] = []
 
     @property
     def supported_features(self) -> CoverEntityFeature:
@@ -688,3 +690,33 @@ async def async_stop_cover(self, **kwargs: Any) -> None:
     async def async_stop_cover_tilt(self, **kwargs: Any) -> None:
         """Stop the cover tilt."""
         await self._device.gateway.covers.stop_cover_tilt(self.info_object)
+
+    def restore_external_state_attributes(
+        self,
+        *,
+        state: Literal[
+            "open", "opening", "closed", "closing"
+        ],  # FIXME: why must these be expanded?
+        target_lift_position: int | None,
+        target_tilt_position: int | None,
+    ):
+        """Restore external state attributes."""
+
+        def refresh_state():
+            refresh_task = asyncio.create_task(
+                self._device.gateway.entities.refresh_state(self.info_object)
+            )
+            self._tasks.append(refresh_task)
+            refresh_task.add_done_callback(self._tasks.remove)
+
+        task = asyncio.create_task(
+            self._device.gateway.covers.restore_external_state_attributes(
+                self.info_object,
+                state=state,
+                target_lift_position=target_lift_position,
+                target_tilt_position=target_tilt_position,
+            )
+        )
+        self._tasks.append(task)
+        task.add_done_callback(self._tasks.remove)
+        task.add_done_callback(lambda _: refresh_state())

zha/application/platforms/cover/websocket_api.py

@@ -155,6 +155,29 @@ async def stop_cover_tilt(
     )
 
 
+class CoverRestoreExternalStateAttributesCommand(PlatformEntityCommand):
+    """Cover restore external state attributes command."""
+
+    command: Literal[APICommands.COVER_RESTORE_EXTERNAL_STATE_ATTRIBUTES] = (
+        APICommands.COVER_RESTORE_EXTERNAL_STATE_ATTRIBUTES
+    )
+    platform: str = Platform.COVER
+    state: Literal["open", "opening", "closed", "closing"]
+    target_lift_position: int
+    target_tilt_position: int
+
+
+@decorators.websocket_command(CoverRestoreExternalStateAttributesCommand)
+@decorators.async_response
+async def restore_cover_external_state_attributes(
+    server: Server, client: Client, command: CoverRestoreExternalStateAttributesCommand
+) -> None:
+    """Stop the cover tilt."""
+    await execute_platform_entity_command(
+        server, client, command, "restore_external_state_attributes"
+    )
+
+
 def load_api(server: Server) -> None:
     """Load the api command handlers."""
     register_api_command(server, open_cover)
@@ -165,3 +188,4 @@ def load_api(server: Server) -> None:
     register_api_command(server, close_cover_tilt)
     register_api_command(server, set_tilt_position)
     register_api_command(server, stop_cover_tilt)
+    register_api_command(server, restore_cover_external_state_attributes)

zha/application/platforms/websocket_api.py

@@ -60,9 +60,14 @@ async def execute_platform_entity_command(
         action = getattr(platform_entity, method_name)
         arg_spec = inspect.getfullargspec(action)
         if arg_spec.varkw:
-            await action(**command.model_dump(exclude_none=True))
+            if inspect.iscoroutinefunction(action):
+                await action(**command.model_dump(exclude_none=True))
+            else:
+                action(**command.model_dump(exclude_none=True))
+        elif inspect.iscoroutinefunction(action):
+            await action()
         else:
-            await action()  # the only argument is self
+            action()  # the only argument is self
 
     except Exception as err:
         _LOGGER.exception("Error executing command: %s", method_name, exc_info=err)

zha/websocket/client/helpers.py

@@ -27,6 +27,7 @@
     CoverCloseTiltCommand,
     CoverOpenCommand,
     CoverOpenTiltCommand,
+    CoverRestoreExternalStateAttributesCommand,
     CoverSetPositionCommand,
     CoverSetTiltPositionCommand,
     CoverStopCommand,
@@ -359,6 +360,24 @@ async def stop_cover_tilt(
         )
         return await self._client.async_send_command(command)
 
+    async def restore_external_state_attributes(
+        self,
+        cover_platform_entity: BasePlatformEntityInfo,
+        state: Literal["open", "opening", "closed", "closing"],
+        target_lift_position: int,
+        target_tilt_position: int,
+    ) -> WebSocketCommandResponse:
+        """Stop a cover tilt."""
+        ensure_platform_entity(cover_platform_entity, Platform.COVER)
+        command = CoverRestoreExternalStateAttributesCommand(
+            ieee=cover_platform_entity.device_ieee,
+            unique_id=cover_platform_entity.unique_id,
+            state=state,
+            target_lift_position=target_lift_position,
+            target_tilt_position=target_tilt_position,
+        )
+        return await self._client.async_send_command(command)
+
 
 class FanHelper:
     """Helper to issue fan commands."""

zha/websocket/const.py

@@ -61,6 +61,7 @@ class APICommands(StrEnum):
     COVER_SET_POSITION = "cover_set_position"
     COVER_SET_TILT_POSITION = "cover_set_tilt_position"
     COVER_STOP_TILT = "cover_stop_tilt"
+    COVER_RESTORE_EXTERNAL_STATE_ATTRIBUTES = "cover_restore_external_state_attributes"
 
     FAN_TURN_ON = "fan_turn_on"
     FAN_TURN_OFF = "fan_turn_off"

zha/websocket/server/api/model.py

@@ -79,6 +79,7 @@ class WebSocketCommand(BaseModel):
         APICommands.COVER_CLOSE_TILT,
         APICommands.COVER_SET_TILT_POSITION,
         APICommands.COVER_STOP_TILT,
+        APICommands.COVER_RESTORE_EXTERNAL_STATE_ATTRIBUTES,
         APICommands.CLIMATE_SET_TEMPERATURE,
         APICommands.CLIMATE_SET_HVAC_MODE,
         APICommands.CLIMATE_SET_FAN_MODE,
@@ -129,6 +130,7 @@ class ErrorResponse(WebSocketCommandResponse):
         "error.cover_set_tilt_position",
         "error.cover_stop",
         "error.cover_stop_tilt",
+        "error.cover_restore_external_state_attributes",
         "error.climate_set_fan_mode",
         "error.climate_set_hvac_mode",
         "error.climate_set_preset_mode",
@@ -182,6 +184,7 @@ class DefaultResponse(WebSocketCommandResponse):
         "cover_open_tilt",
         "cover_close_tilt",
         "cover_set_tilt_position",
+        "cover_restore_external_state_attributes",
         "climate_set_fan_mode",
         "climate_set_hvac_mode",
         "climate_set_preset_mode",