Avoid multiple executor jobs to add wemo devices (#112484)

This commit is contained in:
J. Nick Koston 2024-03-12 07:49:17 -10:00 committed by GitHub
parent 0f414d8ac5
commit 5ae207001f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -6,7 +6,7 @@ import asyncio
from dataclasses import dataclass, fields from dataclasses import dataclass, fields
from datetime import timedelta from datetime import timedelta
import logging import logging
from typing import Literal from typing import TYPE_CHECKING, Literal
from pywemo import Insight, LongPressMixin, WeMoDevice from pywemo import Insight, LongPressMixin, WeMoDevice
from pywemo.exceptions import ActionException, PyWeMoException from pywemo.exceptions import ActionException, PyWeMoException
@ -92,7 +92,7 @@ class DeviceCoordinator(DataUpdateCoordinator[None]): # pylint: disable=hass-en
options: Options | None = None options: Options | None = None
def __init__(self, hass: HomeAssistant, wemo: WeMoDevice, device_id: str) -> None: def __init__(self, hass: HomeAssistant, wemo: WeMoDevice) -> None:
"""Initialize DeviceCoordinator.""" """Initialize DeviceCoordinator."""
super().__init__( super().__init__(
hass, hass,
@ -102,11 +102,16 @@ class DeviceCoordinator(DataUpdateCoordinator[None]): # pylint: disable=hass-en
) )
self.hass = hass self.hass = hass
self.wemo = wemo self.wemo = wemo
self.device_id = device_id self.device_id: str | None = None
self.device_info = _create_device_info(wemo) self.device_info = _create_device_info(wemo)
self.supports_long_press = isinstance(wemo, LongPressMixin) self.supports_long_press = isinstance(wemo, LongPressMixin)
self.update_lock = asyncio.Lock() self.update_lock = asyncio.Lock()
@callback
def async_setup(self, device_id: str) -> None:
"""Set up the device coordinator."""
self.device_id = device_id
def subscription_callback( def subscription_callback(
self, _device: WeMoDevice, event_type: str, params: str self, _device: WeMoDevice, event_type: str, params: str
) -> None: ) -> None:
@ -130,6 +135,9 @@ class DeviceCoordinator(DataUpdateCoordinator[None]): # pylint: disable=hass-en
async def async_shutdown(self) -> None: async def async_shutdown(self) -> None:
"""Unregister push subscriptions and remove from coordinators dict.""" """Unregister push subscriptions and remove from coordinators dict."""
await super().async_shutdown() await super().async_shutdown()
if TYPE_CHECKING:
# mypy doesn't known that the device_id is set in async_setup.
assert self.device_id is not None
del _async_coordinators(self.hass)[self.device_id] del _async_coordinators(self.hass)[self.device_id]
assert self.options # Always set by async_register_device. assert self.options # Always set by async_register_device.
if self.options.enable_subscription: if self.options.enable_subscription:
@ -222,7 +230,10 @@ class DeviceCoordinator(DataUpdateCoordinator[None]): # pylint: disable=hass-en
async def _async_update_data(self) -> None: async def _async_update_data(self) -> None:
"""Update WeMo state.""" """Update WeMo state."""
# No need to poll if the device will push updates. # No need to poll if the device will push updates.
if not self.should_poll: # The device_id will not be set until after the first
# update so we should not check should_poll until after
# the device_id is set.
if self.device_id and not self.should_poll:
return return
# If an update is in progress, we don't do anything. # If an update is in progress, we don't do anything.
@ -266,15 +277,15 @@ async def async_register_device(
hass: HomeAssistant, config_entry: ConfigEntry, wemo: WeMoDevice hass: HomeAssistant, config_entry: ConfigEntry, wemo: WeMoDevice
) -> DeviceCoordinator: ) -> DeviceCoordinator:
"""Register a device with home assistant and enable pywemo event callbacks.""" """Register a device with home assistant and enable pywemo event callbacks."""
# Ensure proper communication with the device and get the initial state. device = DeviceCoordinator(hass, wemo)
await hass.async_add_executor_job(wemo.get_state, True) await device.async_refresh()
if not device.last_update_success and device.last_exception:
raise device.last_exception
device_registry = async_get_device_registry(hass) device_registry = async_get_device_registry(hass)
entry = device_registry.async_get_or_create( entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, **_create_device_info(wemo) config_entry_id=config_entry.entry_id, **_create_device_info(wemo)
) )
device.async_setup(device_id=entry.id)
device = DeviceCoordinator(hass, wemo, entry.id)
_async_coordinators(hass)[entry.id] = device _async_coordinators(hass)[entry.id] = device
config_entry.async_on_unload( config_entry.async_on_unload(