diff --git a/custom_components/detailed_hello_world_push/__init__.py b/custom_components/detailed_hello_world_push/__init__.py
index ae5f619..1b893ce 100644
--- a/custom_components/detailed_hello_world_push/__init__.py
+++ b/custom_components/detailed_hello_world_push/__init__.py
@@ -1,4 +1,5 @@
 """The Detailed Hello World Push integration."""
+
 from __future__ import annotations
 
 from homeassistant.config_entries import ConfigEntry
@@ -6,18 +7,19 @@
 from homeassistant.const import Platform
 
 from . import hub
-from .const import DOMAIN
 
 # List of platforms to support. There should be a matching .py file for each,
 # eg <cover.py> and <sensor.py>
 PLATFORMS = [Platform.SENSOR, Platform.COVER]
 
+type HubConfigEntry = ConfigEntry[hub.Hub]
+
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: HubConfigEntry) -> bool:
     """Set up Hello World from a config entry."""
     # Store an instance of the "connecting" class that does the work of speaking
     # with your actual devices.
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = hub.Hub(hass, entry.data["host"])
+    entry.runtime_data = hub.Hub(hass, entry.data["host"])
 
     # This creates each HA object for each platform your device requires.
     # It's done by calling the `async_setup_entry` function in each platform module.
@@ -31,7 +33,5 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     # needs to unload itself, and remove callbacks. See the classes for further
     # details
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
diff --git a/custom_components/detailed_hello_world_push/cover.py b/custom_components/detailed_hello_world_push/cover.py
index 66ac1b9..32878fc 100644
--- a/custom_components/detailed_hello_world_push/cover.py
+++ b/custom_components/detailed_hello_world_push/cover.py
@@ -10,10 +10,10 @@
     CoverEntityFeature,
     CoverEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import HubConfigEntry
 from .const import DOMAIN
 
 
@@ -21,13 +21,13 @@
 # hass.config_entries.async_forward_entry_setup call)
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HubConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Add cover for passed config_entry in HA."""
-    # The hub is loaded from the associated hass.data entry that was created in the
+    # The hub is loaded from the associated entry runtime data that was set in the
     # __init__.async_setup_entry function
-    hub = hass.data[DOMAIN][config_entry.entry_id]
+    hub = config_entry.runtime_data
 
     # Add all entities to HA
     async_add_entities(HelloWorldCover(roller) for roller in hub.rollers)
diff --git a/custom_components/detailed_hello_world_push/sensor.py b/custom_components/detailed_hello_world_push/sensor.py
index 9825b6f..b45e3d2 100644
--- a/custom_components/detailed_hello_world_push/sensor.py
+++ b/custom_components/detailed_hello_world_push/sensor.py
@@ -9,17 +9,16 @@
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
-    SensorEntity,
-    SensorStateClass,
 )
 from homeassistant.const import (
-    ATTR_VOLTAGE,
     PERCENTAGE,
     LIGHT_LUX,
-    ATTR_UNIT_OF_MEASUREMENT,
 )
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import HubConfigEntry
 from .const import DOMAIN
 
 
@@ -27,9 +26,13 @@
 # Note how both entities for each roller sensor (battry and illuminance) are added at
 # the same time to the same list. This way only a single async_add_devices call is
 # required.
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: HubConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
     """Add sensors for passed config_entry in HA."""
-    hub = hass.data[DOMAIN][config_entry.entry_id]
+    hub = config_entry.runtime_data
 
     new_devices = []
     for roller in hub.rollers: