From 00621ad512e2f377f48cad5c27d8fc2e9529f60b Mon Sep 17 00:00:00 2001
From: G Johansson <goran.johansson@shiftit.se>
Date: Mon, 24 Jun 2024 20:53:49 +0200
Subject: [PATCH] Use runtime data in version (#120363)

---
 homeassistant/components/version/__init__.py      | 13 ++++++-------
 homeassistant/components/version/binary_sensor.py |  9 ++++-----
 homeassistant/components/version/diagnostics.py   |  8 +++-----
 homeassistant/components/version/sensor.py        |  9 ++++-----
 4 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/homeassistant/components/version/__init__.py b/homeassistant/components/version/__init__.py
index 4112cc51e46..cf13821dc8a 100644
--- a/homeassistant/components/version/__init__.py
+++ b/homeassistant/components/version/__init__.py
@@ -16,15 +16,16 @@ from .const import (
     CONF_CHANNEL,
     CONF_IMAGE,
     CONF_SOURCE,
-    DOMAIN,
     PLATFORMS,
 )
 from .coordinator import VersionDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
+type VersionConfigEntry = ConfigEntry[VersionDataUpdateCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: VersionConfigEntry) -> bool:
     """Set up the version integration from a config entry."""
 
     board = entry.data[CONF_BOARD]
@@ -50,14 +51,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: VersionConfigEntry) -> bool:
     """Unload the config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        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/version/binary_sensor.py b/homeassistant/components/version/binary_sensor.py
index ff4f51e409f..827029e1d8c 100644
--- a/homeassistant/components/version/binary_sensor.py
+++ b/homeassistant/components/version/binary_sensor.py
@@ -9,13 +9,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME, EntityCategory, __version__ as HA_VERSION
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import CONF_SOURCE, DEFAULT_NAME, DOMAIN
-from .coordinator import VersionDataUpdateCoordinator
+from . import VersionConfigEntry
+from .const import CONF_SOURCE, DEFAULT_NAME
 from .entity import VersionEntity
 
 HA_VERSION_OBJECT = AwesomeVersion(HA_VERSION)
@@ -23,11 +22,11 @@ HA_VERSION_OBJECT = AwesomeVersion(HA_VERSION)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: VersionConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up version binary_sensors."""
-    coordinator: VersionDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
     if (source := config_entry.data[CONF_SOURCE]) == "local":
         return
 
diff --git a/homeassistant/components/version/diagnostics.py b/homeassistant/components/version/diagnostics.py
index 194027d6ef4..ca7318f468b 100644
--- a/homeassistant/components/version/diagnostics.py
+++ b/homeassistant/components/version/diagnostics.py
@@ -6,20 +6,18 @@ from typing import Any
 
 from attr import asdict
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
-from .const import DOMAIN
-from .coordinator import VersionDataUpdateCoordinator
+from . import VersionConfigEntry
 
 
 async def async_get_config_entry_diagnostics(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: VersionConfigEntry,
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: VersionDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
     device_registry = dr.async_get(hass)
     entity_registry = er.async_get(hass)
 
diff --git a/homeassistant/components/version/sensor.py b/homeassistant/components/version/sensor.py
index 6b0565b8cb3..e1d552bcd36 100644
--- a/homeassistant/components/version/sensor.py
+++ b/homeassistant/components/version/sensor.py
@@ -5,24 +5,23 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from .const import CONF_SOURCE, DEFAULT_NAME, DOMAIN
-from .coordinator import VersionDataUpdateCoordinator
+from . import VersionConfigEntry
+from .const import CONF_SOURCE, DEFAULT_NAME
 from .entity import VersionEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: VersionConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up version sensors."""
-    coordinator: VersionDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     if (entity_name := entry.data[CONF_NAME]) == DEFAULT_NAME:
         entity_name = entry.title