From 5a0e47be48b61cbb63ff734803cdbd7ffd3c3579 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Wed, 23 Oct 2024 18:21:25 +0200
Subject: [PATCH] Use runtime_data in bang_olufsen (#129037)

---
 .../components/bang_olufsen/__init__.py       | 24 ++++++++-----------
 .../components/bang_olufsen/media_player.py   | 12 ++++++----
 tests/components/bang_olufsen/test_init.py    |  3 ++-
 3 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/homeassistant/components/bang_olufsen/__init__.py b/homeassistant/components/bang_olufsen/__init__.py
index e11df6ad5ed..c8ba1f1c3dc 100644
--- a/homeassistant/components/bang_olufsen/__init__.py
+++ b/homeassistant/components/bang_olufsen/__init__.py
@@ -31,10 +31,12 @@ class BangOlufsenData:
     client: MozartClient
 
 
+type BangOlufsenConfigEntry = ConfigEntry[BangOlufsenData]
+
 PLATFORMS = [Platform.MEDIA_PLAYER]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: BangOlufsenConfigEntry) -> bool:
     """Set up from a config entry."""
 
     # Remove casts to str
@@ -67,10 +69,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     websocket = BangOlufsenWebsocket(hass, entry, client)
 
     # Add the websocket and API client
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = BangOlufsenData(
-        websocket,
-        client,
-    )
+    entry.runtime_data = BangOlufsenData(websocket, client)
 
     # Start WebSocket connection
     await client.connect_notifications(remote_control=True, reconnect=True)
@@ -80,15 +79,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, entry: BangOlufsenConfigEntry
+) -> bool:
     """Unload a config entry."""
     # Close the API client and WebSocket notification listener
-    hass.data[DOMAIN][entry.entry_id].client.disconnect_notifications()
-    await hass.data[DOMAIN][entry.entry_id].client.close_api_client()
+    entry.runtime_data.client.disconnect_notifications()
+    await entry.runtime_data.client.close_api_client()
 
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
diff --git a/homeassistant/components/bang_olufsen/media_player.py b/homeassistant/components/bang_olufsen/media_player.py
index ecf571d5456..7c6ea640b38 100644
--- a/homeassistant/components/bang_olufsen/media_player.py
+++ b/homeassistant/components/bang_olufsen/media_player.py
@@ -56,7 +56,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 . import BangOlufsenConfigEntry
 from .const import (
     BANG_OLUFSEN_STATES,
     CONF_BEOLINK_JID,
@@ -96,14 +96,16 @@ BANG_OLUFSEN_FEATURES = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BangOlufsenConfigEntry,
     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, data.client)])
+    async_add_entities(
+        new_entities=[
+            BangOlufsenMediaPlayer(config_entry, config_entry.runtime_data.client)
+        ]
+    )
 
 
 class BangOlufsenMediaPlayer(BangOlufsenEntity, MediaPlayerEntity):
diff --git a/tests/components/bang_olufsen/test_init.py b/tests/components/bang_olufsen/test_init.py
index 3eb98e956be..5b809488ed8 100644
--- a/tests/components/bang_olufsen/test_init.py
+++ b/tests/components/bang_olufsen/test_init.py
@@ -85,6 +85,7 @@ async def test_unload_entry(
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
     assert mock_config_entry.state == ConfigEntryState.LOADED
+    assert hasattr(mock_config_entry, "runtime_data")
 
     # Unload entry
     await hass.config_entries.async_unload(mock_config_entry.entry_id)
@@ -94,5 +95,5 @@ async def test_unload_entry(
     assert mock_mozart_client.close_api_client.call_count == 1
 
     # Ensure that the entry is not loaded and has been removed from hass
-    assert mock_config_entry.entry_id not in hass.data[DOMAIN]
+    assert not hasattr(mock_config_entry, "runtime_data")
     assert mock_config_entry.state == ConfigEntryState.NOT_LOADED