Use runtime_data in fivem (#137632)

Author: epenet

homeassistant/components/fivem/__init__.py

@@ -6,28 +6,26 @@
 
 from fivem import FiveMServerOfflineError
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PORT, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
-from .const import DOMAIN
-from .coordinator import FiveMDataUpdateCoordinator
+from .coordinator import FiveMConfigEntry, FiveMDataUpdateCoordinator
 
 PLATFORMS: list[Platform] = [Platform.BINARY_SENSOR, Platform.SENSOR]
 
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: FiveMConfigEntry) -> bool:
     """Set up FiveM from a config entry."""
     _LOGGER.debug(
         "Create FiveM server instance for '%s:%s'",
         entry.data[CONF_HOST],
         entry.data[CONF_PORT],
     )
 
-    coordinator = FiveMDataUpdateCoordinator(hass, entry.data, entry.entry_id)
+    coordinator = FiveMDataUpdateCoordinator(hass, entry)
 
     try:
         await coordinator.initialize()
@@ -36,16 +34,13 @@ 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: FiveMConfigEntry) -> bool:
     """Unload a 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)

homeassistant/components/fivem/binary_sensor.py

@@ -7,11 +7,11 @@
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, NAME_STATUS
+from .const import NAME_STATUS
+from .coordinator import FiveMConfigEntry
 from .entity import FiveMEntity, FiveMEntityDescription
 
 
@@ -33,11 +33,11 @@ class FiveMBinarySensorEntityDescription(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: FiveMConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FiveM binary sensor platform."""
-    coordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     async_add_entities(
         [FiveMSensorEntity(coordinator, description) for description in BINARY_SENSORS]

homeassistant/components/fivem/coordinator.py

@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
 from datetime import timedelta
 import logging
 from typing import Any
 
 from fivem import FiveM, FiveMServerOfflineError
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PORT
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
@@ -26,26 +26,32 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+type FiveMConfigEntry = ConfigEntry[FiveMDataUpdateCoordinator]
+
 
 class FiveMDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
     """Class to manage fetching FiveM data."""
 
-    def __init__(
-        self, hass: HomeAssistant, config_data: Mapping[str, Any], unique_id: str
-    ) -> None:
+    def __init__(self, hass: HomeAssistant, entry: FiveMConfigEntry) -> None:
         """Initialize server instance."""
-        self.unique_id = unique_id
+        self.unique_id = entry.entry_id
         self.server = None
         self.version = None
         self.game_name: str | None = None
 
-        self.host = config_data[CONF_HOST]
+        self.host = entry.data[CONF_HOST]
 
-        self._fivem = FiveM(self.host, config_data[CONF_PORT])
+        self._fivem = FiveM(self.host, entry.data[CONF_PORT])
 
         update_interval = timedelta(seconds=SCAN_INTERVAL)
 
-        super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=update_interval)
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=entry,
+            name=DOMAIN,
+            update_interval=update_interval,
+        )
 
     async def initialize(self) -> None:
         """Initialize the FiveM server."""

homeassistant/components/fivem/sensor.py

@@ -3,22 +3,21 @@
 from dataclasses import dataclass
 
 from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
 from .const import (
     ATTR_PLAYERS_LIST,
     ATTR_RESOURCES_LIST,
-    DOMAIN,
     NAME_PLAYERS_MAX,
     NAME_PLAYERS_ONLINE,
     NAME_RESOURCES,
     UNIT_PLAYERS_MAX,
     UNIT_PLAYERS_ONLINE,
     UNIT_RESOURCES,
 )
+from .coordinator import FiveMConfigEntry
 from .entity import FiveMEntity, FiveMEntityDescription
 
 
@@ -50,11 +49,11 @@ class FiveMSensorEntityDescription(SensorEntityDescription, FiveMEntityDescripti
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: FiveMConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the FiveM sensor platform."""
-    coordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     # Add sensor entities.
     async_add_entities(