Store transmission coordinator in runtime_data (#119502)

store transmission coordinator in runtime_data
This commit is contained in:
Rami Mosleh 2024-06-12 16:38:35 +03:00 committed by GitHub
parent 2ca580898d
commit e065c70969
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 23 additions and 24 deletions

View File

@ -15,7 +15,7 @@ from transmission_rpc.error import (
) )
import voluptuous as vol import voluptuous as vol
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry, ConfigEntryState
from homeassistant.const import ( from homeassistant.const import (
CONF_HOST, CONF_HOST,
CONF_ID, CONF_ID,
@ -102,8 +102,12 @@ SERVICE_STOP_TORRENT_SCHEMA = vol.All(
) )
) )
type TransmissionConfigEntry = ConfigEntry[TransmissionDataUpdateCoordinator]
async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
async def async_setup_entry(
hass: HomeAssistant, config_entry: TransmissionConfigEntry
) -> bool:
"""Set up the Transmission Component.""" """Set up the Transmission Component."""
@callback @callback
@ -135,7 +139,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
await hass.async_add_executor_job(coordinator.init_torrent_list) await hass.async_add_executor_job(coordinator.init_torrent_list)
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)
@ -204,9 +208,12 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
if unload_ok := await hass.config_entries.async_unload_platforms( if unload_ok := await hass.config_entries.async_unload_platforms(
config_entry, PLATFORMS config_entry, PLATFORMS
): ):
hass.data[DOMAIN].pop(config_entry.entry_id) loaded_entries = [
entry
if not hass.data[DOMAIN]: for entry in hass.config_entries.async_entries(DOMAIN)
if entry.state == ConfigEntryState.LOADED
]
if len(loaded_entries) == 1:
hass.services.async_remove(DOMAIN, SERVICE_ADD_TORRENT) hass.services.async_remove(DOMAIN, SERVICE_ADD_TORRENT)
hass.services.async_remove(DOMAIN, SERVICE_REMOVE_TORRENT) hass.services.async_remove(DOMAIN, SERVICE_REMOVE_TORRENT)
hass.services.async_remove(DOMAIN, SERVICE_START_TORRENT) hass.services.async_remove(DOMAIN, SERVICE_START_TORRENT)

View File

@ -14,7 +14,6 @@ from homeassistant.components.sensor import (
SensorEntity, SensorEntity,
SensorEntityDescription, SensorEntityDescription,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_IDLE, UnitOfDataRate from homeassistant.const import STATE_IDLE, UnitOfDataRate
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
@ -22,6 +21,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.helpers.update_coordinator import CoordinatorEntity
from . import TransmissionConfigEntry
from .const import ( from .const import (
DOMAIN, DOMAIN,
STATE_ATTR_TORRENT_INFO, STATE_ATTR_TORRENT_INFO,
@ -134,14 +134,12 @@ SENSOR_TYPES: tuple[TransmissionSensorEntityDescription, ...] = (
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: TransmissionConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up the Transmission sensors.""" """Set up the Transmission sensors."""
coordinator: TransmissionDataUpdateCoordinator = hass.data[DOMAIN][ coordinator = config_entry.runtime_data
config_entry.entry_id
]
async_add_entities( async_add_entities(
TransmissionSensor(coordinator, description) for description in SENSOR_TYPES TransmissionSensor(coordinator, description) for description in SENSOR_TYPES

View File

@ -2,21 +2,18 @@
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass
import logging
from typing import Any from typing import Any
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
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
from . import TransmissionConfigEntry
from .const import DOMAIN from .const import DOMAIN
from .coordinator import TransmissionDataUpdateCoordinator from .coordinator import TransmissionDataUpdateCoordinator
_LOGGING = logging.getLogger(__name__)
@dataclass(frozen=True, kw_only=True) @dataclass(frozen=True, kw_only=True)
class TransmissionSwitchEntityDescription(SwitchEntityDescription): class TransmissionSwitchEntityDescription(SwitchEntityDescription):
@ -47,14 +44,12 @@ SWITCH_TYPES: tuple[TransmissionSwitchEntityDescription, ...] = (
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: TransmissionConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up the Transmission switch.""" """Set up the Transmission switch."""
coordinator: TransmissionDataUpdateCoordinator = hass.data[DOMAIN][ coordinator = config_entry.runtime_data
config_entry.entry_id
]
async_add_entities( async_add_entities(
TransmissionSwitch(coordinator, description) for description in SWITCH_TYPES TransmissionSwitch(coordinator, description) for description in SWITCH_TYPES

View File

@ -119,7 +119,6 @@ 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 not hass.data[DOMAIN]
@pytest.mark.parametrize( @pytest.mark.parametrize(