From 910c991a58e2718c9d5e81db7603002a9348c902 Mon Sep 17 00:00:00 2001
From: Michael <35783820+mib1185@users.noreply.github.com>
Date: Sat, 4 May 2024 23:39:12 +0200
Subject: [PATCH] Store runtime data inside the config entry in Sun (#116808)

* store runtime data inside the config entry

* move to entry.async_on_unload()
---
 homeassistant/components/sun/__init__.py | 14 +++++++-------
 homeassistant/components/sun/entity.py   |  3 +++
 homeassistant/components/sun/sensor.py   |  7 +++----
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/sun/__init__.py b/homeassistant/components/sun/__init__.py
index 6308594f4bd..8f6f3098ee8 100644
--- a/homeassistant/components/sun/__init__.py
+++ b/homeassistant/components/sun/__init__.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
+from homeassistant.config_entries import SOURCE_IMPORT
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv
@@ -19,7 +19,7 @@ from .const import (  # noqa: F401  # noqa: F401
     STATE_ABOVE_HORIZON,
     STATE_BELOW_HORIZON,
 )
-from .entity import Sun
+from .entity import Sun, SunConfigEntry
 
 CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)
 
@@ -40,19 +40,19 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: SunConfigEntry) -> bool:
     """Set up from a config entry."""
-    hass.data[DOMAIN] = Sun(hass)
+    entry.runtime_data = sun = Sun(hass)
+    entry.async_on_unload(sun.remove_listeners)
     await hass.config_entries.async_forward_entry_setups(entry, [Platform.SENSOR])
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: SunConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(
         entry, [Platform.SENSOR]
     ):
-        sun: Sun = hass.data.pop(DOMAIN)
-        sun.remove_listeners()
+        sun = entry.runtime_data
         hass.states.async_remove(sun.entity_id)
     return unload_ok
diff --git a/homeassistant/components/sun/entity.py b/homeassistant/components/sun/entity.py
index 739784697e0..291f56718a3 100644
--- a/homeassistant/components/sun/entity.py
+++ b/homeassistant/components/sun/entity.py
@@ -8,6 +8,7 @@ from typing import Any
 
 from astral.location import Elevation, Location
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     EVENT_CORE_CONFIG_UPDATE,
     SUN_EVENT_SUNRISE,
@@ -30,6 +31,8 @@ from .const import (
     STATE_BELOW_HORIZON,
 )
 
+SunConfigEntry = ConfigEntry["Sun"]
+
 _LOGGER = logging.getLogger(__name__)
 
 ENTITY_ID = "sun.sun"
diff --git a/homeassistant/components/sun/sensor.py b/homeassistant/components/sun/sensor.py
index 018ba4fa994..e7e621d06cd 100644
--- a/homeassistant/components/sun/sensor.py
+++ b/homeassistant/components/sun/sensor.py
@@ -13,7 +13,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import DEGREE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
@@ -22,7 +21,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
 from .const import DOMAIN, SIGNAL_EVENTS_CHANGED, SIGNAL_POSITION_CHANGED
-from .entity import Sun
+from .entity import Sun, SunConfigEntry
 
 ENTITY_ID_SENSOR_FORMAT = SENSOR_DOMAIN + ".sun_{}"
 
@@ -107,11 +106,11 @@ SENSOR_TYPES: tuple[SunSensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: SunConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up Sun sensor platform."""
 
-    sun: Sun = hass.data[DOMAIN]
+    sun = entry.runtime_data
 
     async_add_entities(
         [SunSensor(sun, description, entry.entry_id) for description in SENSOR_TYPES]