Move API client into dataclass

Fix not all config_flow exceptions caught
Tweaks

homeassistant/components/bangolufsen/__init__.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass
 import logging
 
 from aiohttp.client_exceptions import ClientConnectorError
+from mozart_api.exceptions import ApiException
 from mozart_api.mozart_client import MozartClient
 
 from homeassistant.config_entries import ConfigEntry
@@ -21,29 +22,31 @@ _LOGGER = logging.getLogger(__name__)
 
 @dataclass
 class BangOlufsenData:
-    """Dataclass for storing entities."""
+    """Dataclass for API client and WebSocket client."""
 
     websocket: BangOlufsenWebsocket
+    client: MozartClient
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up from a config entry."""
 
-    client = MozartClient(
-        host=entry.data[CONF_HOST], urllib3_logging_level=logging.DEBUG
-    )
+    client = MozartClient(host=entry.data[CONF_HOST], websocket_reconnect=True)
 
     # Check connection and try to initialize it.
-    async with client:
     try:
         await client.get_battery_state(_request_timeout=3)
-        except (TimeoutError, ClientConnectorError, Exception) as error:
+    except (ApiException, ClientConnectorError, TimeoutError) as error:
+        await client.close()
         raise ConfigEntryNotReady(f"Unable to connect to {entry.title}") from error
 
-    websocket = BangOlufsenWebsocket(hass, entry)
+    websocket = BangOlufsenWebsocket(hass, entry, client)
 
     # Add the websocket
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = BangOlufsenData(websocket)
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = BangOlufsenData(
+        websocket,
+        client,
+    )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     websocket.connect_websocket()
@@ -53,6 +56,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
+    # Close the API client
+    await hass.data[DOMAIN][entry.entry_id].client.close()
+
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
     if unload_ok:

homeassistant/components/bangolufsen/config_flow.py

@@ -6,6 +6,7 @@ import logging
 from typing import Any, TypedDict
 
 from aiohttp.client_exceptions import ClientConnectorError
+from mozart_api.exceptions import ApiException
 from mozart_api.mozart_client import MozartClient
 import voluptuous as vol
 
@@ -100,17 +101,17 @@ class BangOlufsenConfigFlowHandler(ConfigFlow, domain=DOMAIN):
             self._client = MozartClient(self._host, urllib3_logging_level=logging.ERROR)
 
             # Try to get information from Beolink self method.
-            try:
             async with self._client:
+                try:
                     beolink_self = await self._client.get_beolink_self(
                         _request_timeout=3
                     )
-
-            except (TimeoutError, ClientConnectorError) as error:
+                except (ApiException, ClientConnectorError, TimeoutError) as error:
                     return self.async_abort(
                         reason={
-                        TimeoutError: "timeout_error",
+                            ApiException: "api_exception",
                             ClientConnectorError: "client_connector_error",
+                            TimeoutError: "timeout_error",
                         }[type(error)]
                     )
 

homeassistant/components/bangolufsen/entity.py

@@ -1,8 +1,6 @@
 """Entity representing a Bang & Olufsen device."""
 from __future__ import annotations
 
-from collections.abc import Callable
-import logging
 from typing import cast
 
 from mozart_api.models import (
@@ -27,9 +25,12 @@ from .const import DOMAIN
 class BangOlufsenVariables:
     """Shared variables for various classes."""
 
-    def __init__(self, entry: ConfigEntry) -> None:
+    def __init__(self, entry: ConfigEntry, client: MozartClient) -> None:
         """Initialize the object."""
 
+        # Set the MozartClient
+        self._client = client
+
         # get the input from the config entry.
         self.entry: ConfigEntry = entry
 
@@ -40,12 +41,6 @@ class BangOlufsenVariables:
         self._name: str = self.entry.title
         self._unique_id: str = cast(str, self.entry.unique_id)
 
-        self._client: MozartClient = MozartClient(
-            host=self._host,
-            websocket_reconnect=True,
-            urllib3_logging_level=logging.ERROR,
-        )
-
         # Objects that get directly updated by notifications.
         self._playback_metadata: PlaybackContentMetadata = PlaybackContentMetadata()
         self._playback_progress: PlaybackProgress = PlaybackProgress(total_duration=0)
@@ -62,10 +57,10 @@ class BangOlufsenEntity(Entity, BangOlufsenVariables):
 
     _attr_has_entity_name = True
 
-    def __init__(self, entry: ConfigEntry) -> None:
+    def __init__(self, entry: ConfigEntry, client: MozartClient) -> None:
         """Initialize the object."""
-        BangOlufsenVariables.__init__(self, entry)
-        self._dispatchers: list[Callable] = []
+        Entity.__init__(self)
+        BangOlufsenVariables.__init__(self, entry, client)
 
         self._attr_device_info = DeviceInfo(identifiers={(DOMAIN, self._unique_id)})
         self._attr_device_class = None
@@ -77,7 +72,3 @@ class BangOlufsenEntity(Entity, BangOlufsenVariables):
         self._attr_available = connection_state
 
         self.async_write_ha_state()
-
-    async def async_will_remove_from_hass(self) -> None:
-        """Close API client."""
-        await self._client.close()

homeassistant/components/bangolufsen/media_player.py

@@ -29,7 +29,7 @@ from mozart_api.models import (
     VolumeSettings,
     VolumeState,
 )
-from mozart_api.mozart_client import get_highest_resolution_artwork
+from mozart_api.mozart_client import MozartClient, get_highest_resolution_artwork
 
 from homeassistant.components import media_source
 from homeassistant.components.media_player import (
@@ -51,6 +51,7 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.dt import utcnow
 
+from . import BangOlufsenData
 from .const import (
     BANGOLUFSEN_MEDIA_TYPE,
     BANGOLUFSEN_STATES,
@@ -96,10 +97,10 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a Media Player entity from config entry."""
+    data: BangOlufsenData = hass.data[DOMAIN][config_entry.entry_id]
+
     # Add MediaPlayer entity
-    async_add_entities(
-        new_entities=[BangOlufsenMediaPlayer(config_entry)], update_before_add=True
-    )
+    async_add_entities(new_entities=[BangOlufsenMediaPlayer(config_entry, data.client)])
 
 
 class BangOlufsenMediaPlayer(MediaPlayerEntity, BangOlufsenEntity):
@@ -109,9 +110,9 @@ class BangOlufsenMediaPlayer(MediaPlayerEntity, BangOlufsenEntity):
     _attr_icon = "mdi:speaker-wireless"
     _attr_supported_features = BANGOLUFSEN_FEATURES
 
-    def __init__(self, entry: ConfigEntry) -> None:
+    def __init__(self, entry: ConfigEntry, client: MozartClient) -> None:
         """Initialize the media player."""
-        super().__init__(entry)
+        super().__init__(entry, client)
 
         self._beolink_jid: str = self.entry.data[CONF_BEOLINK_JID]
         self._default_volume: int = self.entry.data[CONF_DEFAULT_VOLUME]

homeassistant/components/bangolufsen/strings.json

@@ -1,10 +1,12 @@
 {
   "config": {
     "abort": {
-      "already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]",
       "already_configured": "[%key:common::config_flow::abort::single_instance_allowed%]",
+      "already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]",
+      "api_exception": "[%key:common::config_flow::error::invalid_host%]",
+      "client_connector_error": "[%key:common::config_flow::error::invalid_host%]",
       "timeout_error": "[%key:common::config_flow::error::timeout_connect%]",
-      "client_connector_error": "[%key:common::config_flow::error::invalid_host%]"
+      "value_error": "[%key:common::config_flow::error::invalid_host%]"
     },
     "flow_title": "{name}",
     "step": {

homeassistant/components/bangolufsen/util.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-from typing import cast
-
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceEntry
@@ -17,5 +15,7 @@ def get_device(hass: HomeAssistant | None, unique_id: str) -> DeviceEntry | None
         return None
 
     device_registry = dr.async_get(hass)
-    device = cast(DeviceEntry, device_registry.async_get_device({(DOMAIN, unique_id)}))
+    device = device_registry.async_get_device({(DOMAIN, unique_id)})
+    assert device
+
     return device

homeassistant/components/bangolufsen/websocket.py

@@ -16,6 +16,7 @@ from mozart_api.models import (
     VolumeState,
     WebsocketNotificationTag,
 )
+from mozart_api.mozart_client import MozartClient
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
@@ -36,10 +37,12 @@ _LOGGER = logging.getLogger(__name__)
 class BangOlufsenWebsocket(BangOlufsenVariables):
     """The WebSocket listeners."""
 
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+    def __init__(
+        self, hass: HomeAssistant, entry: ConfigEntry, client: MozartClient
+    ) -> None:
         """Initialize the WebSocket listeners."""
 
-        BangOlufsenVariables.__init__(self, entry)
+        BangOlufsenVariables.__init__(self, entry, client)
 
         self.hass = hass