From 70ae0822811ded8e97740ba20d1be0783f151dcc Mon Sep 17 00:00:00 2001
From: Matthias Alphart <farmio@alphart.net>
Date: Thu, 2 May 2024 13:42:48 +0200
Subject: [PATCH] Use entry.runtime_data in Fronius (#116604)

---
 homeassistant/components/fronius/__init__.py  | 20 +++++++------------
 .../components/fronius/coordinator.py         |  2 +-
 homeassistant/components/fronius/sensor.py    |  7 +++----
 3 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/fronius/__init__.py b/homeassistant/components/fronius/__init__.py
index 1928bb15bc2..c4d1c02ee74 100644
--- a/homeassistant/components/fronius/__init__.py
+++ b/homeassistant/components/fronius/__init__.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
 from datetime import datetime, timedelta
 import logging
 from typing import Final, TypeVar
@@ -42,28 +41,24 @@ PLATFORMS: Final = [Platform.SENSOR]
 
 _FroniusCoordinatorT = TypeVar("_FroniusCoordinatorT", bound=FroniusCoordinatorBase)
 
+FroniusConfigEntry = ConfigEntry["FroniusSolarNet"]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: FroniusConfigEntry) -> bool:
     """Set up fronius from a config entry."""
     host = entry.data[CONF_HOST]
     fronius = Fronius(async_get_clientsession(hass), host)
     solar_net = FroniusSolarNet(hass, entry, fronius)
     await solar_net.init_devices()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = solar_net
+    entry.runtime_data = solar_net
     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: FroniusConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        solar_net = hass.data[DOMAIN].pop(entry.entry_id)
-        while solar_net.cleanup_callbacks:
-            solar_net.cleanup_callbacks.pop()()
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
 async def async_remove_config_entry_device(
@@ -81,7 +76,6 @@ class FroniusSolarNet:
     ) -> None:
         """Initialize FroniusSolarNet class."""
         self.hass = hass
-        self.cleanup_callbacks: list[Callable[[], None]] = []
         self.config_entry = entry
         self.coordinator_lock = asyncio.Lock()
         self.fronius = fronius
@@ -151,7 +145,7 @@ class FroniusSolarNet:
         )
 
         # Setup periodic re-scan
-        self.cleanup_callbacks.append(
+        self.config_entry.async_on_unload(
             async_track_time_interval(
                 self.hass,
                 self._init_devices_inverter,
diff --git a/homeassistant/components/fronius/coordinator.py b/homeassistant/components/fronius/coordinator.py
index 1ecd74a6e09..71ecb4e762e 100644
--- a/homeassistant/components/fronius/coordinator.py
+++ b/homeassistant/components/fronius/coordinator.py
@@ -121,7 +121,7 @@ class FroniusCoordinatorBase(
             async_add_entities(new_entities)
 
         _add_entities_for_unregistered_descriptors()
-        self.solar_net.cleanup_callbacks.append(
+        self.solar_net.config_entry.async_on_unload(
             self.async_add_listener(_add_entities_for_unregistered_descriptors)
         )
 
diff --git a/homeassistant/components/fronius/sensor.py b/homeassistant/components/fronius/sensor.py
index 2d79086d8ba..3b283c33326 100644
--- a/homeassistant/components/fronius/sensor.py
+++ b/homeassistant/components/fronius/sensor.py
@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     POWER_VOLT_AMPERE_REACTIVE,
@@ -44,7 +43,7 @@ from .const import (
 )
 
 if TYPE_CHECKING:
-    from . import FroniusSolarNet
+    from . import FroniusConfigEntry
     from .coordinator import (
         FroniusCoordinatorBase,
         FroniusInverterUpdateCoordinator,
@@ -60,11 +59,11 @@ ENERGY_VOLT_AMPERE_REACTIVE_HOUR: Final = "varh"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: FroniusConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Fronius sensor entities based on a config entry."""
-    solar_net: FroniusSolarNet = hass.data[DOMAIN][config_entry.entry_id]
+    solar_net = config_entry.runtime_data
 
     for inverter_coordinator in solar_net.inverter_coordinators:
         inverter_coordinator.add_entities_for_seen_keys(