Use config entry runtime_data in azure event hub (#127082)

This commit is contained in:
epenet 2024-09-30 10:22:12 +02:00 committed by GitHub
parent 3caf6c0e31
commit 40f808e9be
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 10 deletions

View File

@ -7,6 +7,7 @@ from collections.abc import Callable
from datetime import datetime from datetime import datetime
import json import json
import logging import logging
from types import MappingProxyType
from typing import Any from typing import Any
from azure.eventhub import EventData, EventDataBatch from azure.eventhub import EventData, EventDataBatch
@ -36,12 +37,13 @@ from .const import (
CONF_MAX_DELAY, CONF_MAX_DELAY,
CONF_SEND_INTERVAL, CONF_SEND_INTERVAL,
DATA_FILTER, DATA_FILTER,
DATA_HUB,
DEFAULT_MAX_DELAY, DEFAULT_MAX_DELAY,
DOMAIN, DOMAIN,
FILTER_STATES, FILTER_STATES,
) )
type AzureEventHubConfigEntry = ConfigEntry[AzureEventHub]
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
CONFIG_SCHEMA = vol.Schema( CONFIG_SCHEMA = vol.Schema(
@ -92,7 +94,9 @@ async def async_setup(hass: HomeAssistant, yaml_config: ConfigType) -> bool:
return True return True
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_setup_entry(
hass: HomeAssistant, entry: AzureEventHubConfigEntry
) -> bool:
"""Do the setup based on the config entry and the filter from yaml.""" """Do the setup based on the config entry and the filter from yaml."""
hass.data.setdefault(DOMAIN, {DATA_FILTER: FILTER_SCHEMA({})}) hass.data.setdefault(DOMAIN, {DATA_FILTER: FILTER_SCHEMA({})})
hub = AzureEventHub( hub = AzureEventHub(
@ -104,21 +108,24 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
await hub.async_test_connection() await hub.async_test_connection()
except EventHubError as err: except EventHubError as err:
raise ConfigEntryNotReady("Could not connect to Azure Event Hub") from err raise ConfigEntryNotReady("Could not connect to Azure Event Hub") from err
hass.data[DOMAIN][DATA_HUB] = hub entry.runtime_data = hub
entry.async_on_unload(hub.async_stop)
entry.async_on_unload(entry.add_update_listener(async_update_listener)) entry.async_on_unload(entry.add_update_listener(async_update_listener))
await hub.async_start() await hub.async_start()
return True return True
async def async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None: async def async_update_listener(
hass: HomeAssistant, entry: AzureEventHubConfigEntry
) -> None:
"""Update listener for options.""" """Update listener for options."""
hass.data[DOMAIN][DATA_HUB].update_options(entry.options) entry.runtime_data.update_options(entry.options)
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(
hass: HomeAssistant, entry: AzureEventHubConfigEntry
) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
hub = hass.data[DOMAIN].pop(DATA_HUB)
await hub.async_stop()
return True return True
@ -172,7 +179,7 @@ class AzureEventHub:
await self.async_send(None) await self.async_send(None)
await self._queue.join() await self._queue.join()
def update_options(self, new_options: dict[str, Any]) -> None: def update_options(self, new_options: MappingProxyType[str, Any]) -> None:
"""Update options.""" """Update options."""
self._send_interval = new_options[CONF_SEND_INTERVAL] self._send_interval = new_options[CONF_SEND_INTERVAL]

View File

@ -17,7 +17,6 @@ CONF_EVENT_HUB_CON_STRING = "event_hub_connection_string"
CONF_SEND_INTERVAL = "send_interval" CONF_SEND_INTERVAL = "send_interval"
CONF_MAX_DELAY = "max_delay" CONF_MAX_DELAY = "max_delay"
CONF_FILTER = DATA_FILTER = "filter" CONF_FILTER = DATA_FILTER = "filter"
DATA_HUB = "hub"
STEP_USER = "user" STEP_USER = "user"
STEP_SAS = "sas" STEP_SAS = "sas"