Incorporate various improvements for the ws66i integration

homeassistant/components/ws66i/__init__.py

@@ -94,8 +94,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         zones=zones,
     )
 
+    @callback
     def shutdown(event):
-        """Close the WS66i connection to the amplifier and save snapshots."""
+        """Close the WS66i connection to the amplifier."""
         ws66i.close()
 
     entry.async_on_unload(entry.add_update_listener(_update_listener))
@@ -119,6 +120,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
-async def _update_listener(hass: HomeAssistant, entry: ConfigEntry):
+async def _update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Handle options update."""
     await hass.config_entries.async_reload(entry.entry_id)

homeassistant/components/ws66i/config_flow.py

@@ -1,5 +1,6 @@
 """Config flow for WS66i 6-Zone Amplifier integration."""
 import logging
+from typing import Any
 
 from pyws66i import WS66i, get_ws66i
 import voluptuous as vol
@@ -50,22 +51,34 @@ def _sources_from_config(data):
     }
 
 
-async def validate_input(hass: core.HomeAssistant, input_data):
-    """Validate the user input allows us to connect.
+def _verify_connection(ws66i: WS66i) -> bool:
+    """Verify a connection can be made to the WS66i."""
+    try:
+        ws66i.open()
+    except ConnectionError as err:
+        raise CannotConnect from err
+
+    # Connection successful. Verify correct port was opened
+    # Test on FIRST_ZONE because this zone will always be valid
+    ret_val = ws66i.zone_status(FIRST_ZONE)
+
+    ws66i.close()
+
+    return bool(ret_val)
+
+
+async def validate_input(
+    hass: core.HomeAssistant, input_data: dict[str, Any]
+) -> dict[str, Any]:
+    """Validate the user input.
 
     Data has the keys from DATA_SCHEMA with values provided by the user.
     """
     ws66i: WS66i = get_ws66i(input_data[CONF_IP_ADDRESS])
-    await hass.async_add_executor_job(ws66i.open)
-    # No exception. run a simple test to make sure we opened correct port
-    # Test on FIRST_ZONE because this zone will always be valid
-    ret_val = await hass.async_add_executor_job(ws66i.zone_status, FIRST_ZONE)
-    if ret_val is None:
-        ws66i.close()
-        raise ConnectionError("Not a valid WS66i connection")
 
-    # Validation done. No issues. Close the connection
-    ws66i.close()
+    is_valid: bool = await hass.async_add_executor_job(_verify_connection, ws66i)
+    if not is_valid:
+        raise CannotConnect("Not a valid WS66i connection")
 
     # Return info that you want to store in the config entry.
     return {CONF_IP_ADDRESS: input_data[CONF_IP_ADDRESS]}
@@ -82,17 +95,18 @@ async def async_step_user(self, user_input=None):
         if user_input is not None:
             try:
                 info = await validate_input(self.hass, user_input)
-                # Data is valid. Add default values for options flow.
+            except CannotConnect:
+                errors["base"] = "cannot_connect"
+            except Exception:  # pylint: disable=broad-except
+                _LOGGER.exception("Unexpected exception")
+                errors["base"] = "unknown"
+            else:
+                # Data is valid. Create a config entry.
                 return self.async_create_entry(
                     title="WS66i Amp",
                     data=info,
                     options={CONF_SOURCES: INIT_OPTIONS_DEFAULT},
                 )
-            except ConnectionError:
-                errors["base"] = "cannot_connect"
-            except Exception:  # pylint: disable=broad-except
-                _LOGGER.exception("Unexpected exception")
-                errors["base"] = "unknown"
 
         return self.async_show_form(
             step_id="user", data_schema=DATA_SCHEMA, errors=errors

homeassistant/components/ws66i/const.py

@@ -1,4 +1,5 @@
 """Constants for the Soundavo WS66i 6-Zone Amplifier Media Player component."""
+from datetime import timedelta
 
 DOMAIN = "ws66i"
 
@@ -20,5 +21,6 @@
     "6": "Source 6",
 }
 
-SERVICE_SNAPSHOT = "snapshot"
-SERVICE_RESTORE = "restore"
+POLL_INTERVAL = timedelta(seconds=30)
+
+MAX_VOL = 38

homeassistant/components/ws66i/coordinator.py

@@ -1,20 +1,19 @@
 """Coordinator for WS66i."""
 from __future__ import annotations
 
-from datetime import timedelta
 import logging
 
 from pyws66i import WS66i, ZoneStatus
 
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-_LOGGER = logging.getLogger(__name__)
+from .const import POLL_INTERVAL
 
-POLL_INTERVAL = timedelta(seconds=30)
+_LOGGER = logging.getLogger(__name__)
 
 
-class Ws66iDataUpdateCoordinator(DataUpdateCoordinator):
+class Ws66iDataUpdateCoordinator(DataUpdateCoordinator[list[ZoneStatus]]):
     """DataUpdateCoordinator to gather data for WS66i Zones."""
 
     def __init__(
@@ -43,11 +42,9 @@ def _update_all_zones(self) -> list[ZoneStatus]:
 
             data.append(data_zone)
 
-        # HA will call my entity's _handle_coordinator_update()
         return data
 
     async def _async_update_data(self) -> list[ZoneStatus]:
         """Fetch data for each of the zones."""
-        # HA will call my entity's _handle_coordinator_update()
-        # The data I pass back here can be accessed through coordinator.data.
+        # The data that is returned here can be accessed through coordinator.data.
         return await self.hass.async_add_executor_job(self._update_all_zones)

homeassistant/components/ws66i/media_player.py

@@ -1,6 +1,4 @@
 """Support for interfacing with WS66i 6 zone home audio controller."""
-from copy import deepcopy
-
 from pyws66i import WS66i, ZoneStatus
 
 from homeassistant.components.media_player import (
@@ -10,22 +8,16 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity import DeviceInfo
-from homeassistant.helpers.entity_platform import (
-    AddEntitiesCallback,
-    async_get_current_platform,
-)
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import DOMAIN, SERVICE_RESTORE, SERVICE_SNAPSHOT
+from .const import DOMAIN, MAX_VOL
 from .coordinator import Ws66iDataUpdateCoordinator
 from .models import Ws66iData
 
 PARALLEL_UPDATES = 1
 
-MAX_VOL = 38
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -48,23 +40,8 @@ async def async_setup_entry(
         for idx, zone_id in enumerate(ws66i_data.zones)
     )
 
-    # Set up services
-    platform = async_get_current_platform()
-
-    platform.async_register_entity_service(
-        SERVICE_SNAPSHOT,
-        {},
-        "snapshot",
-    )
 
-    platform.async_register_entity_service(
-        SERVICE_RESTORE,
-        {},
-        "async_restore",
-    )
-
-
-class Ws66iZone(CoordinatorEntity, MediaPlayerEntity):
+class Ws66iZone(CoordinatorEntity[Ws66iDataUpdateCoordinator], MediaPlayerEntity):
     """Representation of a WS66i amplifier zone."""
 
     def __init__(
@@ -82,8 +59,6 @@ def __init__(
         self._ws66i_data: Ws66iData = ws66i_data
         self._zone_id: int = zone_id
         self._zone_id_idx: int = data_idx
-        self._coordinator = coordinator
-        self._snapshot: ZoneStatus = None
         self._status: ZoneStatus = coordinator.data[data_idx]
         self._attr_source_list = ws66i_data.sources.name_list
         self._attr_unique_id = f"{entry_id}_{self._zone_id}"
@@ -131,20 +106,6 @@ def _async_update_attrs_write_ha_state(self) -> None:
         self._set_attrs_from_status()
         self.async_write_ha_state()
 
-    @callback
-    def snapshot(self):
-        """Save zone's current state."""
-        self._snapshot = deepcopy(self._status)
-
-    async def async_restore(self):
-        """Restore saved state."""
-        if not self._snapshot:
-            raise HomeAssistantError("There is no snapshot to restore")
-
-        await self.hass.async_add_executor_job(self._ws66i.restore_zone, self._snapshot)
-        self._status = self._snapshot
-        self._async_update_attrs_write_ha_state()
-
     async def async_select_source(self, source):
         """Set input source."""
         idx = self._ws66i_data.sources.name_id[source]
@@ -180,24 +141,30 @@ async def async_mute_volume(self, mute):
 
     async def async_set_volume_level(self, volume):
         """Set volume level, range 0..1."""
-        await self.hass.async_add_executor_job(
-            self._ws66i.set_volume, self._zone_id, int(volume * MAX_VOL)
-        )
-        self._status.volume = int(volume * MAX_VOL)
+        await self.hass.async_add_executor_job(self._set_volume, int(volume * MAX_VOL))
         self._async_update_attrs_write_ha_state()
 
     async def async_volume_up(self):
         """Volume up the media player."""
         await self.hass.async_add_executor_job(
-            self._ws66i.set_volume, self._zone_id, min(self._status.volume + 1, MAX_VOL)
+            self._set_volume, min(self._status.volume + 1, MAX_VOL)
         )
-        self._status.volume = min(self._status.volume + 1, MAX_VOL)
         self._async_update_attrs_write_ha_state()
 
     async def async_volume_down(self):
         """Volume down media player."""
         await self.hass.async_add_executor_job(
-            self._ws66i.set_volume, self._zone_id, max(self._status.volume - 1, 0)
+            self._set_volume, max(self._status.volume - 1, 0)
         )
-        self._status.volume = max(self._status.volume - 1, 0)
         self._async_update_attrs_write_ha_state()
+
+    def _set_volume(self, volume: int) -> None:
+        """Set the volume of the media player."""
+        # Can't set a new volume level when this zone is muted.
+        # Follow behavior of keypads, where zone is unmuted when volume changes.
+        if self._status.mute:
+            self._ws66i.set_mute(self._zone_id, False)
+            self._status.mute = False
+
+        self._ws66i.set_volume(self._zone_id, volume)
+        self._status.volume = volume

homeassistant/components/ws66i/models.py

@@ -7,8 +7,6 @@
 
 from .coordinator import Ws66iDataUpdateCoordinator
 
-# A dataclass is basically a struct in C/C++
-
 
 @dataclass
 class SourceRep:

homeassistant/components/ws66i/strings.json

@@ -11,9 +11,6 @@
     "error": {
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
       "unknown": "[%key:common::config_flow::error::unknown%]"
-    },
-    "abort": {
-      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
     }
   },
   "options": {

homeassistant/components/ws66i/translations/en.json

@@ -1,8 +1,5 @@
 {
     "config": {
-        "abort": {
-            "already_configured": "Device is already configured"
-        },
         "error": {
             "cannot_connect": "Failed to connect",
             "unknown": "Unexpected error"

tests/components/ws66i/test_config_flow.py

@@ -1,7 +1,7 @@
 """Test the WS66i 6-Zone Amplifier config flow."""
 from unittest.mock import patch
 
-from homeassistant import config_entries, data_entry_flow, setup
+from homeassistant import config_entries, data_entry_flow
 from homeassistant.components.ws66i.const import (
     CONF_SOURCE_1,
     CONF_SOURCE_2,
@@ -15,15 +15,15 @@
 )
 from homeassistant.const import CONF_IP_ADDRESS
 
+from .test_media_player import AttrDict
+
 from tests.common import MockConfigEntry
-from tests.components.ws66i.test_media_player import AttrDict
 
 CONFIG = {CONF_IP_ADDRESS: "1.1.1.1"}
 
 
 async def test_form(hass):
     """Test we get the form."""
-    await setup.async_setup_component(hass, "persistent_notification", {})
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )