Store imap coordinator in runtime_data (#119611)
This commit is contained in:
parent
50fe29ccc1
commit
bb2883a5a8
3 changed files with 22 additions and 32 deletions
|
@ -4,12 +4,11 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from aioimaplib import IMAP4_SSL, AioImapException, Response
|
||||
import voluptuous as vol
|
||||
|
||||
from homeassistant.config_entries import ConfigEntry
|
||||
from homeassistant.config_entries import ConfigEntry, ConfigEntryState
|
||||
from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
|
||||
from homeassistant.core import (
|
||||
HomeAssistant,
|
||||
|
@ -29,6 +28,7 @@ from homeassistant.helpers.typing import ConfigType
|
|||
|
||||
from .const import CONF_ENABLE_PUSH, DOMAIN
|
||||
from .coordinator import (
|
||||
ImapDataUpdateCoordinator,
|
||||
ImapMessage,
|
||||
ImapPollingDataUpdateCoordinator,
|
||||
ImapPushDataUpdateCoordinator,
|
||||
|
@ -65,17 +65,18 @@ SERVICE_MOVE_SCHEMA = _SERVICE_UID_SCHEMA.extend(
|
|||
SERVICE_DELETE_SCHEMA = _SERVICE_UID_SCHEMA
|
||||
SERVICE_FETCH_TEXT_SCHEMA = _SERVICE_UID_SCHEMA
|
||||
|
||||
type ImapConfigEntry = ConfigEntry[ImapDataUpdateCoordinator]
|
||||
|
||||
|
||||
async def async_get_imap_client(hass: HomeAssistant, entry_id: str) -> IMAP4_SSL:
|
||||
"""Get IMAP client and connect."""
|
||||
if hass.data[DOMAIN].get(entry_id) is None:
|
||||
if (entry := hass.config_entries.async_get_entry(entry_id)) is None or (
|
||||
entry.state is not ConfigEntryState.LOADED
|
||||
):
|
||||
raise ServiceValidationError(
|
||||
translation_domain=DOMAIN,
|
||||
translation_key="invalid_entry",
|
||||
)
|
||||
entry = hass.config_entries.async_get_entry(entry_id)
|
||||
if TYPE_CHECKING:
|
||||
assert entry is not None
|
||||
try:
|
||||
client = await connect_to_server(entry.data)
|
||||
except InvalidAuth as exc:
|
||||
|
@ -235,7 +236,7 @@ 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: ImapConfigEntry) -> bool:
|
||||
"""Set up imap from a config entry."""
|
||||
try:
|
||||
imap_client: IMAP4_SSL = await connect_to_server(dict(entry.data))
|
||||
|
@ -255,12 +256,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
|||
else:
|
||||
coordinator_class = ImapPollingDataUpdateCoordinator
|
||||
|
||||
coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = (
|
||||
coordinator_class(hass, imap_client, entry)
|
||||
)
|
||||
coordinator: ImapDataUpdateCoordinator = coordinator_class(hass, imap_client, entry)
|
||||
await coordinator.async_config_entry_first_refresh()
|
||||
|
||||
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
|
||||
entry.runtime_data = coordinator
|
||||
|
||||
entry.async_on_unload(
|
||||
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, coordinator.shutdown)
|
||||
|
@ -271,11 +270,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||
async def async_unload_entry(hass: HomeAssistant, entry: ImapConfigEntry) -> bool:
|
||||
"""Unload a config entry."""
|
||||
if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
|
||||
coordinator: (
|
||||
ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator
|
||||
) = hass.data[DOMAIN].pop(entry.entry_id)
|
||||
coordinator = entry.runtime_data
|
||||
await coordinator.shutdown()
|
||||
return unload_ok
|
||||
|
|
|
@ -5,18 +5,16 @@ from __future__ import annotations
|
|||
from typing import Any
|
||||
|
||||
from homeassistant.components.diagnostics import async_redact_data
|
||||
from homeassistant.config_entries import ConfigEntry
|
||||
from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
|
||||
from homeassistant.core import HomeAssistant, callback
|
||||
|
||||
from .const import DOMAIN
|
||||
from .coordinator import ImapDataUpdateCoordinator
|
||||
from . import ImapConfigEntry
|
||||
|
||||
REDACT_CONFIG = {CONF_PASSWORD, CONF_USERNAME}
|
||||
|
||||
|
||||
async def async_get_config_entry_diagnostics(
|
||||
hass: HomeAssistant, entry: ConfigEntry
|
||||
hass: HomeAssistant, entry: ImapConfigEntry
|
||||
) -> dict[str, Any]:
|
||||
"""Return diagnostics for a config entry."""
|
||||
return _async_get_diagnostics(hass, entry)
|
||||
|
@ -25,11 +23,11 @@ async def async_get_config_entry_diagnostics(
|
|||
@callback
|
||||
def _async_get_diagnostics(
|
||||
hass: HomeAssistant,
|
||||
entry: ConfigEntry,
|
||||
entry: ImapConfigEntry,
|
||||
) -> dict[str, Any]:
|
||||
"""Return diagnostics for a config entry."""
|
||||
redacted_config = async_redact_data(entry.data, REDACT_CONFIG)
|
||||
coordinator: ImapDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
|
||||
coordinator = entry.runtime_data
|
||||
|
||||
return {
|
||||
"config": redacted_config,
|
||||
|
|
|
@ -7,15 +7,15 @@ from homeassistant.components.sensor import (
|
|||
SensorEntityDescription,
|
||||
SensorStateClass,
|
||||
)
|
||||
from homeassistant.config_entries import ConfigEntry
|
||||
from homeassistant.const import CONF_USERNAME
|
||||
from homeassistant.core import HomeAssistant
|
||||
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
|
||||
from homeassistant.helpers.entity_platform import AddEntitiesCallback
|
||||
from homeassistant.helpers.update_coordinator import CoordinatorEntity
|
||||
|
||||
from . import ImapPollingDataUpdateCoordinator, ImapPushDataUpdateCoordinator
|
||||
from . import ImapConfigEntry
|
||||
from .const import DOMAIN
|
||||
from .coordinator import ImapDataUpdateCoordinator
|
||||
|
||||
IMAP_MAIL_COUNT_DESCRIPTION = SensorEntityDescription(
|
||||
key="imap_mail_count",
|
||||
|
@ -27,27 +27,22 @@ IMAP_MAIL_COUNT_DESCRIPTION = SensorEntityDescription(
|
|||
|
||||
|
||||
async def async_setup_entry(
|
||||
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
|
||||
hass: HomeAssistant, entry: ImapConfigEntry, async_add_entities: AddEntitiesCallback
|
||||
) -> None:
|
||||
"""Set up the Imap sensor."""
|
||||
|
||||
coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator = (
|
||||
hass.data[DOMAIN][entry.entry_id]
|
||||
)
|
||||
coordinator = entry.runtime_data
|
||||
async_add_entities([ImapSensor(coordinator, IMAP_MAIL_COUNT_DESCRIPTION)])
|
||||
|
||||
|
||||
class ImapSensor(
|
||||
CoordinatorEntity[ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator],
|
||||
SensorEntity,
|
||||
):
|
||||
class ImapSensor(CoordinatorEntity[ImapDataUpdateCoordinator], SensorEntity):
|
||||
"""Representation of an IMAP sensor."""
|
||||
|
||||
_attr_has_entity_name = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
coordinator: ImapPushDataUpdateCoordinator | ImapPollingDataUpdateCoordinator,
|
||||
coordinator: ImapDataUpdateCoordinator,
|
||||
description: SensorEntityDescription,
|
||||
) -> None:
|
||||
"""Initialize the sensor."""
|
||||
|
|
Loading…
Add table
Reference in a new issue