Store imap coordinator in runtime_data (#119611)

This commit is contained in:
Rami Mosleh 2024-06-13 17:58:05 +03:00 committed by GitHub
parent 50fe29ccc1
commit bb2883a5a8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 22 additions and 32 deletions

View file

@ -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

View file

@ -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,

View file

@ -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."""