Explicitly pass in the config_entry in roborock coordinator (#137970)

explicitly pass in the config_entry in coordinator
This commit is contained in:
Michael 2025-02-09 05:53:44 +01:00 committed by GitHub
parent 6e84280e3c
commit a526baa831
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 77 additions and 47 deletions

View File

@ -4,7 +4,6 @@ from __future__ import annotations
import asyncio
from collections.abc import Coroutine
from dataclasses import dataclass
from datetime import timedelta
import logging
from typing import Any
@ -21,35 +20,23 @@ from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
from roborock.version_a01_apis import RoborockMqttClientA01
from roborock.web_api import RoborockApiClient
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_USERNAME, EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from .const import CONF_BASE_URL, CONF_USER_DATA, DOMAIN, PLATFORMS
from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01
from .coordinator import (
RoborockConfigEntry,
RoborockCoordinators,
RoborockDataUpdateCoordinator,
RoborockDataUpdateCoordinatorA01,
)
from .roborock_storage import async_remove_map_storage
SCAN_INTERVAL = timedelta(seconds=30)
_LOGGER = logging.getLogger(__name__)
type RoborockConfigEntry = ConfigEntry[RoborockCoordinators]
@dataclass
class RoborockCoordinators:
"""Roborock coordinators type."""
v1: list[RoborockDataUpdateCoordinator]
a01: list[RoborockDataUpdateCoordinatorA01]
def values(
self,
) -> list[RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01]:
"""Return all coordinators."""
return self.v1 + self.a01
async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> bool:
"""Set up roborock from a config entry."""
@ -95,7 +82,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
# Get a Coordinator if the device is available or if we have connected to the device before
coordinators = await asyncio.gather(
*build_setup_functions(
hass, device_map, user_data, product_info, home_data.rooms
hass, entry, device_map, user_data, product_info, home_data.rooms
),
return_exceptions=True,
)
@ -142,6 +129,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
def build_setup_functions(
hass: HomeAssistant,
entry: RoborockConfigEntry,
device_map: dict[str, HomeDataDevice],
user_data: UserData,
product_info: dict[str, HomeDataProduct],
@ -156,7 +144,12 @@ def build_setup_functions(
"""Create a list of setup functions that can later be called asynchronously."""
return [
setup_device(
hass, user_data, device, product_info[device.product_id], home_data_rooms
hass,
entry,
user_data,
device,
product_info[device.product_id],
home_data_rooms,
)
for device in device_map.values()
]
@ -164,6 +157,7 @@ def build_setup_functions(
async def setup_device(
hass: HomeAssistant,
entry: RoborockConfigEntry,
user_data: UserData,
device: HomeDataDevice,
product_info: HomeDataProduct,
@ -172,10 +166,10 @@ async def setup_device(
"""Set up a coordinator for a given device."""
if device.pv == "1.0":
return await setup_device_v1(
hass, user_data, device, product_info, home_data_rooms
hass, entry, user_data, device, product_info, home_data_rooms
)
if device.pv == "A01":
return await setup_device_a01(hass, user_data, device, product_info)
return await setup_device_a01(hass, entry, user_data, device, product_info)
_LOGGER.warning(
"Not adding device %s because its protocol version %s or category %s is not supported",
device.duid,
@ -187,6 +181,7 @@ async def setup_device(
async def setup_device_v1(
hass: HomeAssistant,
entry: RoborockConfigEntry,
user_data: UserData,
device: HomeDataDevice,
product_info: HomeDataProduct,
@ -212,7 +207,7 @@ async def setup_device_v1(
await mqtt_client.async_release()
raise
coordinator = RoborockDataUpdateCoordinator(
hass, device, networking, product_info, mqtt_client, home_data_rooms
hass, entry, device, networking, product_info, mqtt_client, home_data_rooms
)
try:
await coordinator.async_config_entry_first_refresh()
@ -246,6 +241,7 @@ async def setup_device_v1(
async def setup_device_a01(
hass: HomeAssistant,
entry: RoborockConfigEntry,
user_data: UserData,
device: HomeDataDevice,
product_info: HomeDataProduct,
@ -254,7 +250,9 @@ async def setup_device_a01(
mqtt_client = RoborockMqttClientA01(
user_data, DeviceData(device, product_info.name), product_info.category
)
coord = RoborockDataUpdateCoordinatorA01(hass, device, product_info, mqtt_client)
coord = RoborockDataUpdateCoordinatorA01(
hass, entry, device, product_info, mqtt_client
)
await coord.async_config_entry_first_refresh()
return coord

View File

@ -16,8 +16,7 @@ from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import RoborockConfigEntry
from .coordinator import RoborockDataUpdateCoordinator
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockCoordinatedEntityV1

View File

@ -11,8 +11,7 @@ from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import RoborockConfigEntry
from .coordinator import RoborockDataUpdateCoordinator
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockEntityV1

View File

@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
from dataclasses import dataclass
from datetime import timedelta
import logging
@ -35,14 +36,32 @@ SCAN_INTERVAL = timedelta(seconds=30)
_LOGGER = logging.getLogger(__name__)
@dataclass
class RoborockCoordinators:
"""Roborock coordinators type."""
v1: list[RoborockDataUpdateCoordinator]
a01: list[RoborockDataUpdateCoordinatorA01]
def values(
self,
) -> list[RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01]:
"""Return all coordinators."""
return self.v1 + self.a01
type RoborockConfigEntry = ConfigEntry[RoborockCoordinators]
class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
"""Class to manage fetching data from the API."""
config_entry: ConfigEntry
config_entry: RoborockConfigEntry
def __init__(
self,
hass: HomeAssistant,
config_entry: RoborockConfigEntry,
device: HomeDataDevice,
device_networking: NetworkInfo,
product_info: HomeDataProduct,
@ -50,7 +69,13 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
home_data_rooms: list[HomeDataRoom],
) -> None:
"""Initialize."""
super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL)
super().__init__(
hass,
_LOGGER,
config_entry=config_entry,
name=DOMAIN,
update_interval=SCAN_INTERVAL,
)
self.roborock_device_info = RoborockHassDeviceInfo(
device,
device_networking,
@ -186,15 +211,24 @@ class RoborockDataUpdateCoordinatorA01(
):
"""Class to manage fetching data from the API for A01 devices."""
config_entry: RoborockConfigEntry
def __init__(
self,
hass: HomeAssistant,
config_entry: RoborockConfigEntry,
device: HomeDataDevice,
product_info: HomeDataProduct,
api: RoborockClientA01,
) -> None:
"""Initialize."""
super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL)
super().__init__(
hass,
_LOGGER,
config_entry=config_entry,
name=DOMAIN,
update_interval=SCAN_INTERVAL,
)
self.api = api
self.device_info = DeviceInfo(
name=device.name,

View File

@ -8,7 +8,7 @@ from homeassistant.components.diagnostics import async_redact_data
from homeassistant.const import CONF_UNIQUE_ID
from homeassistant.core import HomeAssistant
from . import RoborockConfigEntry
from .coordinator import RoborockConfigEntry
TO_REDACT_CONFIG = ["token", "sn", "rruid", CONF_UNIQUE_ID, "username", "uid"]

View File

@ -19,7 +19,6 @@ from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.util import dt as dt_util
from . import RoborockConfigEntry
from .const import (
DEFAULT_DRAWABLES,
DOMAIN,
@ -28,7 +27,7 @@ from .const import (
MAP_FILE_FORMAT,
MAP_SLEEP,
)
from .coordinator import RoborockDataUpdateCoordinator
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockCoordinatedEntityV1

View File

@ -16,8 +16,8 @@ from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import DOMAIN, RoborockConfigEntry
from .coordinator import RoborockDataUpdateCoordinator
from .const import DOMAIN
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockEntityV1
_LOGGER = logging.getLogger(__name__)

View File

@ -13,9 +13,8 @@ from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import RoborockConfigEntry
from .const import MAP_SLEEP
from .coordinator import RoborockDataUpdateCoordinator
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockCoordinatedEntityV1

View File

@ -31,8 +31,11 @@ from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType
from . import RoborockConfigEntry
from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01
from .coordinator import (
RoborockConfigEntry,
RoborockDataUpdateCoordinator,
RoborockDataUpdateCoordinatorA01,
)
from .entity import RoborockCoordinatedEntityA01, RoborockCoordinatedEntityV1

View File

@ -18,8 +18,8 @@ from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import DOMAIN, RoborockConfigEntry
from .coordinator import RoborockDataUpdateCoordinator
from .const import DOMAIN
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockEntityV1
_LOGGER = logging.getLogger(__name__)

View File

@ -18,8 +18,8 @@ from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import DOMAIN, RoborockConfigEntry
from .coordinator import RoborockDataUpdateCoordinator
from .const import DOMAIN
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockEntityV1
_LOGGER = logging.getLogger(__name__)

View File

@ -18,14 +18,13 @@ from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import config_validation as cv, entity_platform
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import RoborockConfigEntry
from .const import (
DOMAIN,
GET_MAPS_SERVICE_NAME,
GET_VACUUM_CURRENT_POSITION_SERVICE_NAME,
SET_VACUUM_GOTO_POSITION_SERVICE_NAME,
)
from .coordinator import RoborockDataUpdateCoordinator
from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
from .entity import RoborockCoordinatedEntityV1
from .image import ColorsPalette, ImageConfig, RoborockMapDataParser, Sizes