Store Mikrotik coordinator in runtime_data (#119594)

This commit is contained in:
Rami Mosleh 2024-06-13 12:53:32 +03:00 committed by GitHub
parent f5b86154b4
commit 030fe6df4a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 10 additions and 16 deletions

View File

@ -14,8 +14,12 @@ CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False)
PLATFORMS = [Platform.DEVICE_TRACKER] PLATFORMS = [Platform.DEVICE_TRACKER]
type MikrotikConfigEntry = ConfigEntry[MikrotikDataUpdateCoordinator]
async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
async def async_setup_entry(
hass: HomeAssistant, config_entry: MikrotikConfigEntry
) -> bool:
"""Set up the Mikrotik component.""" """Set up the Mikrotik component."""
try: try:
api = await hass.async_add_executor_job(get_api, dict(config_entry.data)) api = await hass.async_add_executor_job(get_api, dict(config_entry.data))
@ -28,7 +32,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
await hass.async_add_executor_job(coordinator.api.get_hub_details) await hass.async_add_executor_job(coordinator.api.get_hub_details)
await coordinator.async_config_entry_first_refresh() await coordinator.async_config_entry_first_refresh()
hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = coordinator config_entry.runtime_data = coordinator
await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS) await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
@ -47,9 +51,4 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
if unload_ok := await hass.config_entries.async_unload_platforms( return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
config_entry, PLATFORMS
):
hass.data[DOMAIN].pop(config_entry.entry_id)
return unload_ok

View File

@ -9,26 +9,23 @@ from homeassistant.components.device_tracker import (
ScannerEntity, ScannerEntity,
SourceType, SourceType,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import entity_registry as er from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.helpers.update_coordinator import CoordinatorEntity
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from .const import DOMAIN from . import MikrotikConfigEntry
from .coordinator import Device, MikrotikDataUpdateCoordinator from .coordinator import Device, MikrotikDataUpdateCoordinator
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: MikrotikConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up device tracker for Mikrotik component.""" """Set up device tracker for Mikrotik component."""
coordinator: MikrotikDataUpdateCoordinator = hass.data[DOMAIN][ coordinator = config_entry.runtime_data
config_entry.entry_id
]
tracked: dict[str, MikrotikDataUpdateCoordinatorTracker] = {} tracked: dict[str, MikrotikDataUpdateCoordinatorTracker] = {}

View File

@ -6,7 +6,6 @@ from librouteros.exceptions import ConnectionClosed, LibRouterosError
import pytest import pytest
from homeassistant.components import mikrotik from homeassistant.components import mikrotik
from homeassistant.components.mikrotik.const import DOMAIN
from homeassistant.config_entries import ConfigEntryState from homeassistant.config_entries import ConfigEntryState
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
@ -84,4 +83,3 @@ async def test_unload_entry(hass: HomeAssistant) -> None:
await hass.async_block_till_done() await hass.async_block_till_done()
assert entry.state is ConfigEntryState.NOT_LOADED assert entry.state is ConfigEntryState.NOT_LOADED
assert entry.entry_id not in hass.data[DOMAIN]