Fix WebSocket listener not being awaited

Fix WebSocket remote listener not being started
Rework WebSocket connection initialization to ensure all entities get notification data
Re-add MDNS discovery confirm step
Add IPv4 IP address checking
Improve config flow error messages

custom_components/bangolufsen/__init__.py

@@ -1,6 +1,7 @@
 """The Bang & Olufsen integration."""
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass
 
 from aiohttp.client_exceptions import ClientConnectorError
@@ -11,7 +12,7 @@
 from homeassistant.const import CONF_HOST, CONF_MODEL, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DOMAIN
@@ -20,10 +21,12 @@
 
 @dataclass
 class BangOlufsenData:
-    """Dataclass for API client and coordinator containing WebSocket client."""
+    """Dataclass for API client, coordinator containing WebSocket client and WebSocket initialization variables."""
 
     coordinator: DataUpdateCoordinator
     client: MozartClient
+    entities_initialized: int = 0
+    platforms_initialized: int = 0
 
 
 PLATFORMS = [
@@ -38,15 +41,38 @@ class BangOlufsenData:
 ]
 
 
+async def _start_websocket_listener(
+    hass: HomeAssistant, entry: ConfigEntry, data: BangOlufsenData
+) -> None:
+    """Start WebSocket listener when all entities have been initialized."""
+    entity_registry = er.async_get(hass)
+
+    while True:
+        # Get all entries for entities and filter all disabled entities
+        entries = er.async_entries_for_config_entry(entity_registry, entry.entry_id)
+        expected_entities = len(
+            [entry for entry in entries if not entry.disabled or not entry.disabled_by]
+        )
+
+        # Check if all entities and platforms have been initialized and start WebSocket listener
+        if (
+            expected_entities == data.entities_initialized
+            and len(PLATFORMS) == data.platforms_initialized
+        ):
+            await data.client.connect_notifications(remote_control=True)
+            return
+
+        await asyncio.sleep(0)
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up from a config entry."""
     # Remove casts to str
     assert entry.unique_id
 
+    # Create device now as BangOlufsenWebsocket needs a device for debug logging, firing events etc.
+    # And in order to ensure entity platforms (button, binary_sensor) have device name before the primary (media_player) is initialized
     device_registry = dr.async_get(hass)
-
-    # Create device in order to ensure entity platforms (button, binary_sensor)
-    # have device name before the primary (media_player) is initialized
     device_registry.async_get_or_create(
         config_entry_id=entry.entry_id,
         identifiers={(DOMAIN, entry.unique_id)},
@@ -56,25 +82,35 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     client = MozartClient(host=entry.data[CONF_HOST], websocket_reconnect=True)
 
-    # Check connection and try to initialize it.
+    # Check API connection and try to initialize it.
     try:
         await client.get_battery_state(_request_timeout=3)
     except (ApiException, ClientConnectorError, TimeoutError) as error:
         await client.close_api_client()
         raise ConfigEntryNotReady(f"Unable to connect to {entry.title}") from error
 
+    # Check WebSocket connection
+    if not await client.check_websocket_connection():
+        raise ConfigEntryNotReady(
+            f"Unable to connect to {entry.title} WebSocket notification channel"
+        )
+
+    # Initialize coordinator
     coordinator = BangOlufsenCoordinator(hass, entry, client)
     await coordinator.async_config_entry_first_refresh()
 
+    # Add the websocket and API client
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = BangOlufsenData(
-        coordinator=coordinator,
-        client=client,
+        coordinator, client
     )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
-    client.connect_notifications()
 
-    entry.async_on_unload(entry.add_update_listener(update_listener))
+    # Start WebSocket connection when all entities have been initialized
+    hass.async_create_background_task(
+        _start_websocket_listener(hass, entry, hass.data[DOMAIN][entry.entry_id]),
+        f"{DOMAIN}-{entry.unique_id}-websocket_starter",
+    )
 
     return True
 
@@ -91,8 +127,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
-
-
-async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
-    """Handle options update."""
-    await hass.config_entries.async_reload(entry.entry_id)

custom_components/bangolufsen/binary_sensor.py

@@ -23,6 +23,7 @@
     WEBSOCKET_NOTIFICATION,
 )
 from .entity import BangOlufsenEntity
+from .util import set_platform_initialized
 
 
 async def async_setup_entry(
@@ -32,7 +33,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up Binary Sensor entities from config entry."""
     data: BangOlufsenData = hass.data[DOMAIN][config_entry.entry_id]
-    entities: list[BangOlufsenBinarySensor] = []
+    entities: list[BangOlufsenEntity] = []
 
     # Check if device has a battery
     battery_state = await data.client.get_battery_state()
@@ -48,6 +49,8 @@ async def async_setup_entry(
 
     async_add_entities(new_entities=entities)
 
+    set_platform_initialized(data)
+
 
 class BangOlufsenBinarySensor(BangOlufsenEntity, BinarySensorEntity):
     """Base Binary Sensor class."""
@@ -89,6 +92,8 @@ async def async_added_to_hass(self) -> None:
             )
         )
 
+        self.set_entity_initialized()
+
     async def _update_battery_charging(self, data: BatteryState) -> None:
         """Update battery charging."""
         self._attr_is_on = data.is_charging
@@ -124,6 +129,8 @@ async def async_added_to_hass(self) -> None:
             )
         )
 
+        self.set_entity_initialized()
+
     async def _update_proximity(self, data: WebsocketNotificationTag) -> None:
         """Update proximity."""
         if data.value:

custom_components/bangolufsen/button.py

@@ -17,6 +17,7 @@
 from . import BangOlufsenData
 from .const import CONNECTION_STATUS, DOMAIN, SOURCE_ENUM
 from .entity import BangOlufsenEntity
+from .util import set_platform_initialized
 
 
 async def async_setup_entry(
@@ -26,7 +27,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up Button entities from config entry."""
     data: BangOlufsenData = hass.data[DOMAIN][config_entry.entry_id]
-    entities: list[BangOlufsenButton] = []
+    entities: list[BangOlufsenEntity] = []
 
     # Get available favourites from coordinator.
     favourites = data.coordinator.data.favourites
@@ -40,6 +41,8 @@ async def async_setup_entry(
 
     async_add_entities(new_entities=entities)
 
+    set_platform_initialized(data)
+
 
 class BangOlufsenButton(ButtonEntity, BangOlufsenEntity):
     """Base Button class."""
@@ -87,6 +90,8 @@ async def async_added_to_hass(self) -> None:
 
         self._attr_extra_state_attributes = self._generate_favourite_attributes()
 
+        self.set_entity_initialized()
+
     async def async_press(self) -> None:
         """Handle the action."""
         await self._client.activate_preset(id=self._favourite_id)

custom_components/bangolufsen/config_flow.py

@@ -1,7 +1,7 @@
 """Config flow for the Bang & Olufsen integration."""
 from __future__ import annotations
 
-import ipaddress
+from ipaddress import AddressValueError, IPv4Address
 from typing import Any, TypedDict
 
 from aiohttp.client_exceptions import ClientConnectorError
@@ -21,6 +21,7 @@
     ATTR_SERIAL_NUMBER,
     ATTR_TYPE_NUMBER,
     COMPATIBLE_MODELS,
+    CONF_SERIAL_NUMBER,
     DEFAULT_MODEL,
     DOMAIN,
 )
@@ -35,6 +36,15 @@ class EntryData(TypedDict, total=False):
     name: str
 
 
+# Map exception types to strings
+_exception_map = {
+    ApiException: "api_exception",
+    ClientConnectorError: "client_connector_error",
+    TimeoutError: "timeout_error",
+    AddressValueError: "invalid_ip",
+}
+
+
 class BangOlufsenConfigFlowHandler(ConfigFlow, domain=DOMAIN):
     """Handle a config flow."""
 
@@ -67,14 +77,14 @@ async def async_step_user(
             self._host = user_input[CONF_HOST]
             self._model = user_input[CONF_MODEL]
 
-            # Check if the IP address is a valid address.
+            # Check if the IP address is a valid IPv4 address.
             try:
-                ipaddress.ip_address(self._host)
-            except ValueError:
+                IPv4Address(self._host)
+            except AddressValueError as error:
                 return self.async_show_form(
                     step_id="user",
                     data_schema=data_schema,
-                    errors={"base": "value_error"},
+                    errors={"base": _exception_map[type(error)]},
                 )
 
             self._client = MozartClient(self._host)
@@ -85,7 +95,6 @@ async def async_step_user(
                     beolink_self = await self._client.get_beolink_self(
                         _request_timeout=3
                     )
-
                 except (
                     ApiException,
                     ClientConnectorError,
@@ -94,18 +103,15 @@ async def async_step_user(
                     return self.async_show_form(
                         step_id="user",
                         data_schema=data_schema,
-                        errors={
-                            "base": {
-                                ApiException: "api_exception",
-                                ClientConnectorError: "client_connector_error",
-                                TimeoutError: "timeout_error",
-                            }[type(error)]
-                        },
+                        errors={"base": _exception_map[type(error)]},
                     )
 
             self._beolink_jid = beolink_self.jid
             self._serial_number = beolink_self.jid.split(".")[2].split("@")[0]
 
+            await self.async_set_unique_id(self._serial_number)
+            self._abort_if_unique_id_configured()
+
             return await self._create_entry()
 
         return self.async_show_form(
@@ -122,18 +128,29 @@ async def async_step_zeroconf(
         if ATTR_FRIENDLY_NAME not in discovery_info.properties:
             return self.async_abort(reason="not_mozart_device")
 
+        # Ensure that an IPv4 address is received
         self._host = discovery_info.host
+        try:
+            IPv4Address(self._host)
+        except AddressValueError:
+            return self.async_abort(reason="ipv6_address")
+
         self._model = discovery_info.hostname[:-16].replace("-", " ")
         self._serial_number = discovery_info.properties[ATTR_SERIAL_NUMBER]
         self._beolink_jid = f"{discovery_info.properties[ATTR_TYPE_NUMBER]}.{discovery_info.properties[ATTR_ITEM_NUMBER]}.{self._serial_number}@products.bang-olufsen.com"
 
-        return await self._create_entry()
-
-    async def _create_entry(self) -> FlowResult:
-        """Create the config entry for a discovered or manually configured Bang & Olufsen device."""
         await self.async_set_unique_id(self._serial_number)
         self._abort_if_unique_id_configured(updates={CONF_HOST: self._host})
 
+        # Set the discovered device title
+        self.context["title_placeholders"] = {
+            "name": discovery_info.properties[ATTR_FRIENDLY_NAME]
+        }
+
+        return await self.async_step_zeroconf_confirm()
+
+    async def _create_entry(self) -> FlowResult:
+        """Create the config entry for a discovered or manually configured Bang & Olufsen device."""
         # Ensure that created entities have a unique and easily identifiable id and not a "friendly name"
         self._name = f"{self._model}-{self._serial_number}"
 
@@ -146,3 +163,22 @@ async def _create_entry(self) -> FlowResult:
                 name=self._name,
             ),
         )
+
+    async def async_step_zeroconf_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Confirm the configuration of the device."""
+        if user_input is not None:
+            return await self._create_entry()
+
+        self._set_confirm_only()
+
+        return self.async_show_form(
+            step_id="zeroconf_confirm",
+            description_placeholders={
+                CONF_HOST: self._host,
+                CONF_MODEL: self._model,
+                CONF_SERIAL_NUMBER: self._serial_number,
+            },
+            last_step=True,
+        )