From fa815234be97fa48c43992023a5b0e243a442611 Mon Sep 17 00:00:00 2001
From: Robert Svensson <Kane610@users.noreply.github.com>
Date: Tue, 14 May 2024 21:04:26 +0200
Subject: [PATCH] Make UniFi use runtime data (#117457)
---
homeassistant/components/unifi/__init__.py | 31 +++++++-------
homeassistant/components/unifi/button.py | 12 ++----
homeassistant/components/unifi/config_flow.py | 11 ++---
.../components/unifi/device_tracker.py | 10 ++---
homeassistant/components/unifi/diagnostics.py | 8 ++--
homeassistant/components/unifi/hub/hub.py | 22 +++++-----
homeassistant/components/unifi/image.py | 11 ++---
homeassistant/components/unifi/sensor.py | 6 +--
homeassistant/components/unifi/services.py | 16 +++----
homeassistant/components/unifi/switch.py | 10 ++---
homeassistant/components/unifi/update.py | 7 ++--
tests/common.py | 6 ++-
tests/components/unifi/conftest.py | 7 ++--
tests/components/unifi/test_config_flow.py | 21 +---------
tests/components/unifi/test_hub.py | 11 ++---
tests/components/unifi/test_init.py | 19 ---------
tests/components/unifi/test_services.py | 42 +------------------
17 files changed, 77 insertions(+), 173 deletions(-)
diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py
index 69a6ec423ae..af14bffb8e8 100644
--- a/homeassistant/components/unifi/__init__.py
+++ b/homeassistant/components/unifi/__init__.py
@@ -14,7 +14,9 @@ from homeassistant.helpers.typing import ConfigType
from .const import DOMAIN as UNIFI_DOMAIN, PLATFORMS, UNIFI_WIRELESS_CLIENTS
from .errors import AuthenticationRequired, CannotConnect
from .hub import UnifiHub, get_unifi_api
-from .services import async_setup_services, async_unload_services
+from .services import async_setup_services
+
+UnifiConfigEntry = ConfigEntry[UnifiHub]
SAVE_DELAY = 10
STORAGE_KEY = "unifi_data"
@@ -25,13 +27,17 @@ CONFIG_SCHEMA = cv.config_entry_only_config_schema(UNIFI_DOMAIN)
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Integration doesn't support configuration through configuration.yaml."""
+ async_setup_services(hass)
+
hass.data[UNIFI_WIRELESS_CLIENTS] = wireless_clients = UnifiWirelessClients(hass)
await wireless_clients.async_load()
return True
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+ hass: HomeAssistant, config_entry: UnifiConfigEntry
+) -> bool:
"""Set up the UniFi Network integration."""
hass.data.setdefault(UNIFI_DOMAIN, {})
@@ -44,17 +50,13 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
except AuthenticationRequired as err:
raise ConfigEntryAuthFailed from err
- hub = UnifiHub(hass, config_entry, api)
+ hub = config_entry.runtime_data = UnifiHub(hass, config_entry, api)
await hub.initialize()
- hass.data[UNIFI_DOMAIN][config_entry.entry_id] = hub
await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
hub.async_update_device_registry()
hub.entity_loader.load_entities()
- if len(hass.data[UNIFI_DOMAIN]) == 1:
- async_setup_services(hass)
-
hub.websocket.start()
config_entry.async_on_unload(
@@ -64,21 +66,18 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
return True
-async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+ hass: HomeAssistant, config_entry: UnifiConfigEntry
+) -> bool:
"""Unload a config entry."""
- hub: UnifiHub = hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id)
-
- if not hass.data[UNIFI_DOMAIN]:
- async_unload_services(hass)
-
- return await hub.async_reset()
+ return await config_entry.runtime_data.async_reset()
async def async_remove_config_entry_device(
- hass: HomeAssistant, config_entry: ConfigEntry, device_entry: DeviceEntry
+ hass: HomeAssistant, config_entry: UnifiConfigEntry, device_entry: DeviceEntry
) -> bool:
"""Remove config entry from a device."""
- hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
return not any(
identifier
for _, identifier in device_entry.connections
diff --git a/homeassistant/components/unifi/button.py b/homeassistant/components/unifi/button.py
index 86c38a5bf3d..6684e33e532 100644
--- a/homeassistant/components/unifi/button.py
+++ b/homeassistant/components/unifi/button.py
@@ -29,11 +29,11 @@ from homeassistant.components.button import (
ButtonEntity,
ButtonEntityDescription,
)
-from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from . import UnifiConfigEntry
from .entity import (
HandlerT,
UnifiEntity,
@@ -43,7 +43,6 @@ from .entity import (
async_wlan_available_fn,
async_wlan_device_info_fn,
)
-from .hub import UnifiHub
async def async_restart_device_control_fn(
@@ -123,15 +122,12 @@ ENTITY_DESCRIPTIONS: tuple[UnifiButtonEntityDescription, ...] = (
async def async_setup_entry(
hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: UnifiConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up button platform for UniFi Network integration."""
- UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
- async_add_entities,
- UnifiButtonEntity,
- ENTITY_DESCRIPTIONS,
- requires_admin=True,
+ config_entry.runtime_data.entity_loader.register_platform(
+ async_add_entities, UnifiButtonEntity, ENTITY_DESCRIPTIONS, requires_admin=True
)
diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py
index 79b5e035f41..e703f393d68 100644
--- a/homeassistant/components/unifi/config_flow.py
+++ b/homeassistant/components/unifi/config_flow.py
@@ -36,6 +36,7 @@ from homeassistant.core import HomeAssistant, callback
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device_registry import format_mac
+from . import UnifiConfigEntry
from .const import (
CONF_ALLOW_BANDWIDTH_SENSORS,
CONF_ALLOW_UPTIME_SENSORS,
@@ -163,9 +164,7 @@ class UnifiFlowHandler(ConfigFlow, domain=UNIFI_DOMAIN):
abort_reason = "reauth_successful"
if config_entry:
- hub: UnifiHub | None = self.hass.data.get(UNIFI_DOMAIN, {}).get(
- config_entry.entry_id
- )
+ hub = config_entry.runtime_data
if hub and hub.available:
return self.async_abort(reason="already_configured")
@@ -249,7 +248,7 @@ class UnifiOptionsFlowHandler(OptionsFlow):
hub: UnifiHub
- def __init__(self, config_entry: ConfigEntry) -> None:
+ def __init__(self, config_entry: UnifiConfigEntry) -> None:
"""Initialize UniFi Network options flow."""
self.config_entry = config_entry
self.options = dict(config_entry.options)
@@ -258,9 +257,7 @@ class UnifiOptionsFlowHandler(OptionsFlow):
self, user_input: dict[str, Any] | None = None
) -> ConfigFlowResult:
"""Manage the UniFi Network options."""
- if self.config_entry.entry_id not in self.hass.data[UNIFI_DOMAIN]:
- return self.async_abort(reason="integration_not_setup")
- self.hub = self.hass.data[UNIFI_DOMAIN][self.config_entry.entry_id]
+ self.hub = self.config_entry.runtime_data
self.options[CONF_BLOCK_CLIENT] = self.hub.config.option_block_clients
if self.show_advanced_options:
diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py
index dc48b9c31fe..a1014bfd184 100644
--- a/homeassistant/components/unifi/device_tracker.py
+++ b/homeassistant/components/unifi/device_tracker.py
@@ -18,13 +18,13 @@ from aiounifi.models.device import Device
from aiounifi.models.event import Event, EventKey
from homeassistant.components.device_tracker import DOMAIN, ScannerEntity, SourceType
-from homeassistant.config_entries import ConfigEntry
from homeassistant.core import Event as core_Event, HomeAssistant, callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.helpers.entity_registry as er
import homeassistant.util.dt as dt_util
+from . import UnifiConfigEntry
from .const import DOMAIN as UNIFI_DOMAIN
from .entity import (
HandlerT,
@@ -185,12 +185,12 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = (
@callback
-def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
+def async_update_unique_id(hass: HomeAssistant, config_entry: UnifiConfigEntry) -> None:
"""Normalize client unique ID to have a prefix rather than suffix.
Introduced with release 2023.12.
"""
- hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
ent_reg = er.async_get(hass)
@callback
@@ -210,12 +210,12 @@ def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> No
async def async_setup_entry(
hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: UnifiConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up device tracker for UniFi Network integration."""
async_update_unique_id(hass, config_entry)
- UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
+ config_entry.runtime_data.entity_loader.register_platform(
async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS
)
diff --git a/homeassistant/components/unifi/diagnostics.py b/homeassistant/components/unifi/diagnostics.py
index 7df082ca0a4..21174342594 100644
--- a/homeassistant/components/unifi/diagnostics.py
+++ b/homeassistant/components/unifi/diagnostics.py
@@ -7,13 +7,11 @@ from itertools import chain
from typing import Any
from homeassistant.components.diagnostics import REDACTED, async_redact_data
-from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import format_mac
-from .const import DOMAIN as UNIFI_DOMAIN
-from .hub import UnifiHub
+from . import UnifiConfigEntry
TO_REDACT = {CONF_PASSWORD}
REDACT_CONFIG = {CONF_HOST, CONF_PASSWORD, CONF_USERNAME}
@@ -73,10 +71,10 @@ def async_replace_list_data(
async def async_get_config_entry_diagnostics(
- hass: HomeAssistant, config_entry: ConfigEntry
+ hass: HomeAssistant, config_entry: UnifiConfigEntry
) -> dict[str, Any]:
"""Return diagnostics for a config entry."""
- hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
diag: dict[str, Any] = {}
macs_to_redact: dict[str, str] = {}
diff --git a/homeassistant/components/unifi/hub/hub.py b/homeassistant/components/unifi/hub/hub.py
index f8c1f2517a2..c7615714764 100644
--- a/homeassistant/components/unifi/hub/hub.py
+++ b/homeassistant/components/unifi/hub/hub.py
@@ -3,10 +3,10 @@
from __future__ import annotations
from datetime import datetime
+from typing import TYPE_CHECKING
import aiounifi
-from homeassistant.config_entries import ConfigEntry
from homeassistant.core import Event, HomeAssistant, callback
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.device_registry import (
@@ -22,12 +22,18 @@ from .entity_helper import UnifiEntityHelper
from .entity_loader import UnifiEntityLoader
from .websocket import UnifiWebsocket
+if TYPE_CHECKING:
+ from .. import UnifiConfigEntry
+
class UnifiHub:
"""Manages a single UniFi Network instance."""
def __init__(
- self, hass: HomeAssistant, config_entry: ConfigEntry, api: aiounifi.Controller
+ self,
+ hass: HomeAssistant,
+ config_entry: UnifiConfigEntry,
+ api: aiounifi.Controller,
) -> None:
"""Initialize the system."""
self.hass = hass
@@ -40,13 +46,6 @@ class UnifiHub:
self.site = config_entry.data[CONF_SITE_ID]
self.is_admin = False
- @callback
- @staticmethod
- def get_hub(hass: HomeAssistant, config_entry: ConfigEntry) -> UnifiHub:
- """Get UniFi hub from config entry."""
- hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
- return hub
-
@property
def available(self) -> bool:
"""Websocket connection state."""
@@ -122,15 +121,14 @@ class UnifiHub:
@staticmethod
async def async_config_entry_updated(
- hass: HomeAssistant, config_entry: ConfigEntry
+ hass: HomeAssistant, config_entry: UnifiConfigEntry
) -> None:
"""Handle signals of config entry being updated.
If config entry is updated due to reauth flow
the entry might already have been reset and thus is not available.
"""
- if not (hub := hass.data[UNIFI_DOMAIN].get(config_entry.entry_id)):
- return
+ hub = config_entry.runtime_data
hub.config = UnifiConfig.from_config_entry(config_entry)
async_dispatcher_send(hass, hub.signal_options_update)
diff --git a/homeassistant/components/unifi/image.py b/homeassistant/components/unifi/image.py
index 285477fe133..bbc20e2b06b 100644
--- a/homeassistant/components/unifi/image.py
+++ b/homeassistant/components/unifi/image.py
@@ -14,12 +14,12 @@ from aiounifi.models.api import ApiItemT
from aiounifi.models.wlan import Wlan
from homeassistant.components.image import ImageEntity, ImageEntityDescription
-from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.util.dt as dt_util
+from . import UnifiConfigEntry
from .entity import (
HandlerT,
UnifiEntity,
@@ -65,15 +65,12 @@ ENTITY_DESCRIPTIONS: tuple[UnifiImageEntityDescription, ...] = (
async def async_setup_entry(
hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: UnifiConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up image platform for UniFi Network integration."""
- UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
- async_add_entities,
- UnifiImageEntity,
- ENTITY_DESCRIPTIONS,
- requires_admin=True,
+ config_entry.runtime_data.entity_loader.register_platform(
+ async_add_entities, UnifiImageEntity, ENTITY_DESCRIPTIONS, requires_admin=True
)
diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py
index 2685f075cd5..3fd179f5676 100644
--- a/homeassistant/components/unifi/sensor.py
+++ b/homeassistant/components/unifi/sensor.py
@@ -32,7 +32,6 @@ from homeassistant.components.sensor import (
SensorStateClass,
UnitOfTemperature,
)
-from homeassistant.config_entries import ConfigEntry
from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfDataRate, UnitOfPower
from homeassistant.core import Event as core_Event, HomeAssistant, callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -40,6 +39,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType
import homeassistant.util.dt as dt_util
+from . import UnifiConfigEntry
from .const import DEVICE_STATES
from .entity import (
HandlerT,
@@ -420,11 +420,11 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
async def async_setup_entry(
hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: UnifiConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up sensors for UniFi Network integration."""
- UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
+ config_entry.runtime_data.entity_loader.register_platform(
async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS
)
diff --git a/homeassistant/components/unifi/services.py b/homeassistant/components/unifi/services.py
index 096f4f27dae..5dcc0e9719c 100644
--- a/homeassistant/components/unifi/services.py
+++ b/homeassistant/components/unifi/services.py
@@ -49,13 +49,6 @@ def async_setup_services(hass: HomeAssistant) -> None:
)
-@callback
-def async_unload_services(hass: HomeAssistant) -> None:
- """Unload UniFi Network services."""
- for service in SUPPORTED_SERVICES:
- hass.services.async_remove(UNIFI_DOMAIN, service)
-
-
async def async_reconnect_client(hass: HomeAssistant, data: Mapping[str, Any]) -> None:
"""Try to get wireless client to reconnect to Wi-Fi."""
device_registry = dr.async_get(hass)
@@ -73,9 +66,10 @@ async def async_reconnect_client(hass: HomeAssistant, data: Mapping[str, Any]) -
if mac == "":
return
- for hub in hass.data[UNIFI_DOMAIN].values():
+ for entry in hass.config_entries.async_entries(UNIFI_DOMAIN):
if (
- not hub.available
+ (hub := entry.runtime_data)
+ and not hub.available
or (client := hub.api.clients.get(mac)) is None
or client.is_wired
):
@@ -91,8 +85,8 @@ async def async_remove_clients(hass: HomeAssistant, data: Mapping[str, Any]) ->
- Total time between first seen and last seen is less than 15 minutes.
- Neither IP, hostname nor name is configured.
"""
- for hub in hass.data[UNIFI_DOMAIN].values():
- if not hub.available:
+ for entry in hass.config_entries.async_entries(UNIFI_DOMAIN):
+ if (hub := entry.runtime_data) and not hub.available:
continue
clients_to_remove = []
diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py
index 45357dd67d2..be475803f7e 100644
--- a/homeassistant/components/unifi/switch.py
+++ b/homeassistant/components/unifi/switch.py
@@ -38,13 +38,13 @@ from homeassistant.components.switch import (
SwitchEntity,
SwitchEntityDescription,
)
-from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.helpers.entity_registry as er
+from . import UnifiConfigEntry
from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
from .entity import (
HandlerT,
@@ -270,12 +270,12 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
@callback
-def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
+def async_update_unique_id(hass: HomeAssistant, config_entry: UnifiConfigEntry) -> None:
"""Normalize switch unique ID to have a prefix rather than midfix.
Introduced with release 2023.12.
"""
- hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
ent_reg = er.async_get(hass)
@callback
@@ -299,12 +299,12 @@ def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> No
async def async_setup_entry(
hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: UnifiConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up switches for UniFi Network integration."""
async_update_unique_id(hass, config_entry)
- UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
+ config_entry.runtime_data.entity_loader.register_platform(
async_add_entities,
UnifiSwitchEntity,
ENTITY_DESCRIPTIONS,
diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py
index a8fe3c83427..b3cfc6f1c66 100644
--- a/homeassistant/components/unifi/update.py
+++ b/homeassistant/components/unifi/update.py
@@ -18,17 +18,16 @@ from homeassistant.components.update import (
UpdateEntityDescription,
UpdateEntityFeature,
)
-from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from . import UnifiConfigEntry
from .entity import (
UnifiEntity,
UnifiEntityDescription,
async_device_available_fn,
async_device_device_info_fn,
)
-from .hub import UnifiHub
LOGGER = logging.getLogger(__name__)
@@ -68,11 +67,11 @@ ENTITY_DESCRIPTIONS: tuple[UnifiUpdateEntityDescription, ...] = (
async def async_setup_entry(
hass: HomeAssistant,
- config_entry: ConfigEntry,
+ config_entry: UnifiConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up update entities for UniFi Network integration."""
- UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
+ config_entry.runtime_data.entity_loader.register_platform(
async_add_entities,
UnifiDeviceUpdateEntity,
ENTITY_DESCRIPTIONS,
diff --git a/tests/common.py b/tests/common.py
index b25d730a8cd..55c448fdad2 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -38,7 +38,7 @@ from homeassistant.components.device_automation import ( # noqa: F401
_async_get_device_automation_capabilities as async_get_device_automation_capabilities,
)
from homeassistant.config import async_process_component_config
-from homeassistant.config_entries import ConfigEntry, ConfigFlow
+from homeassistant.config_entries import ConfigEntry, ConfigFlow, _DataT
from homeassistant.const import (
DEVICE_DEFAULT_NAME,
EVENT_HOMEASSISTANT_CLOSE,
@@ -973,9 +973,11 @@ class MockToggleEntity(entity.ToggleEntity):
return None
-class MockConfigEntry(config_entries.ConfigEntry):
+class MockConfigEntry(config_entries.ConfigEntry[_DataT]):
"""Helper for creating config entries that adds some defaults."""
+ runtime_data: _DataT
+
def __init__(
self,
*,
diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py
index 1ef8948ec51..938c26b1730 100644
--- a/tests/components/unifi/conftest.py
+++ b/tests/components/unifi/conftest.py
@@ -9,7 +9,6 @@ from unittest.mock import patch
from aiounifi.models.message import MessageKey
import pytest
-from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
from homeassistant.components.unifi.hub.websocket import RETRY_TIMER
from homeassistant.const import CONTENT_TYPE_JSON
from homeassistant.core import HomeAssistant
@@ -44,7 +43,9 @@ class WebsocketStateManager(asyncio.Event):
Mock api calls done by 'await self.api.login'.
Fail will make 'await self.api.start_websocket' return immediately.
"""
- hub = self.hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
+ hub = self.hass.config_entries.async_get_entry(
+ DEFAULT_CONFIG_ENTRY_ID
+ ).runtime_data
self.aioclient_mock.get(
f"https://{hub.config.host}:1234", status=302
) # Check UniFi OS
@@ -80,7 +81,7 @@ def mock_unifi_websocket(hass):
data: list[dict] | dict | None = None,
):
"""Generate a websocket call."""
- hub = hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
+ hub = hass.config_entries.async_get_entry(DEFAULT_CONFIG_ENTRY_ID).runtime_data
if data and not message:
hub.api.messages.handler(data)
elif data and message:
diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py
index b269392f707..06ada29f911 100644
--- a/tests/components/unifi/test_config_flow.py
+++ b/tests/components/unifi/test_config_flow.py
@@ -278,15 +278,11 @@ async def test_flow_aborts_configuration_updated(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
) -> None:
"""Test config flow aborts since a connected config entry already exists."""
- entry = MockConfigEntry(
- domain=UNIFI_DOMAIN, data={"host": "1.2.3.4", "site": "office"}, unique_id="2"
- )
- entry.add_to_hass(hass)
-
entry = MockConfigEntry(
domain=UNIFI_DOMAIN, data={"host": "1.2.3.4", "site": "site_id"}, unique_id="1"
)
entry.add_to_hass(hass)
+ entry.runtime_data = None
result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -393,7 +389,7 @@ async def test_reauth_flow_update_configuration(
) -> None:
"""Verify reauth flow can update hub configuration."""
config_entry = await setup_unifi_integration(hass, aioclient_mock)
- hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
hub.websocket.available = False
result = await hass.config_entries.flow.async_init(
@@ -572,19 +568,6 @@ async def test_simple_option_flow(
}
-async def test_option_flow_integration_not_setup(
- hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
-) -> None:
- """Test advanced config flow options."""
- config_entry = await setup_unifi_integration(hass, aioclient_mock)
-
- hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id)
- result = await hass.config_entries.options.async_init(config_entry.entry_id)
-
- assert result["type"] is FlowResultType.ABORT
- assert result["reason"] == "integration_not_setup"
-
-
async def test_form_ssdp(hass: HomeAssistant) -> None:
"""Test we get the form with ssdp source."""
diff --git a/tests/components/unifi/test_hub.py b/tests/components/unifi/test_hub.py
index 1fddb623930..579c39daa4f 100644
--- a/tests/components/unifi/test_hub.py
+++ b/tests/components/unifi/test_hub.py
@@ -235,9 +235,6 @@ async def setup_unifi_integration(
await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done()
- if config_entry.entry_id not in hass.data[UNIFI_DOMAIN]:
- return None
-
return config_entry
@@ -254,7 +251,7 @@ async def test_hub_setup(
config_entry = await setup_unifi_integration(
hass, aioclient_mock, system_information_response=SYSTEM_INFORMATION
)
- hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
entry = hub.config.entry
assert len(forward_entry_setup.mock_calls) == 1
@@ -333,7 +330,7 @@ async def test_config_entry_updated(
) -> None:
"""Calling reset when the entry has been setup."""
config_entry = await setup_unifi_integration(hass, aioclient_mock)
- hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
event_call = Mock()
unsub = async_dispatcher_connect(hass, hub.signal_options_update, event_call)
@@ -356,7 +353,7 @@ async def test_reset_after_successful_setup(
) -> None:
"""Calling reset when the entry has been setup."""
config_entry = await setup_unifi_integration(hass, aioclient_mock)
- hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
result = await hub.async_reset()
await hass.async_block_till_done()
@@ -369,7 +366,7 @@ async def test_reset_fails(
) -> None:
"""Calling reset when the entry has been setup can return false."""
config_entry = await setup_unifi_integration(hass, aioclient_mock)
- hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
with patch(
"homeassistant.config_entries.ConfigEntries.async_forward_entry_unload",
diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py
index f358c03d98d..323211272e7 100644
--- a/tests/components/unifi/test_init.py
+++ b/tests/components/unifi/test_init.py
@@ -31,14 +31,6 @@ async def test_setup_with_no_config(hass: HomeAssistant) -> None:
assert UNIFI_DOMAIN not in hass.data
-async def test_successful_config_entry(
- hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
-) -> None:
- """Test that configured options for a host are loaded via config entry."""
- await setup_unifi_integration(hass, aioclient_mock)
- assert hass.data[UNIFI_DOMAIN]
-
-
async def test_setup_entry_fails_config_entry_not_ready(hass: HomeAssistant) -> None:
"""Failed authentication trigger a reauthentication flow."""
with patch(
@@ -65,17 +57,6 @@ async def test_setup_entry_fails_trigger_reauth_flow(hass: HomeAssistant) -> Non
assert hass.data[UNIFI_DOMAIN] == {}
-async def test_unload_entry(
- hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
-) -> None:
- """Test being able to unload an entry."""
- config_entry = await setup_unifi_integration(hass, aioclient_mock)
- assert hass.data[UNIFI_DOMAIN]
-
- assert await hass.config_entries.async_unload(config_entry.entry_id)
- assert not hass.data[UNIFI_DOMAIN]
-
-
async def test_wireless_clients(
hass: HomeAssistant,
hass_storage: dict[str, Any],
diff --git a/tests/components/unifi/test_services.py b/tests/components/unifi/test_services.py
index 3f7da7a63ae..8cd029b1cf5 100644
--- a/tests/components/unifi/test_services.py
+++ b/tests/components/unifi/test_services.py
@@ -1,12 +1,9 @@
"""deCONZ service tests."""
-from unittest.mock import patch
-
from homeassistant.components.unifi.const import CONF_SITE_ID, DOMAIN as UNIFI_DOMAIN
from homeassistant.components.unifi.services import (
SERVICE_RECONNECT_CLIENT,
SERVICE_REMOVE_CLIENTS,
- SUPPORTED_SERVICES,
)
from homeassistant.const import ATTR_DEVICE_ID, CONF_HOST
from homeassistant.core import HomeAssistant
@@ -17,41 +14,6 @@ from .test_hub import setup_unifi_integration
from tests.test_util.aiohttp import AiohttpClientMocker
-async def test_service_setup_and_unload(
- hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
-) -> None:
- """Verify service setup works."""
- config_entry = await setup_unifi_integration(hass, aioclient_mock)
- for service in SUPPORTED_SERVICES:
- assert hass.services.has_service(UNIFI_DOMAIN, service)
-
- assert await hass.config_entries.async_unload(config_entry.entry_id)
- for service in SUPPORTED_SERVICES:
- assert not hass.services.has_service(UNIFI_DOMAIN, service)
-
-
-@patch("homeassistant.core.ServiceRegistry.async_remove")
-@patch("homeassistant.core.ServiceRegistry.async_register")
-async def test_service_setup_and_unload_not_called_if_multiple_integrations_detected(
- register_service_mock,
- remove_service_mock,
- hass: HomeAssistant,
- aioclient_mock: AiohttpClientMocker,
-) -> None:
- """Make sure that services are only setup and removed once."""
- config_entry = await setup_unifi_integration(hass, aioclient_mock)
- register_service_mock.reset_mock()
- config_entry_2 = await setup_unifi_integration(
- hass, aioclient_mock, config_entry_id=2
- )
- register_service_mock.assert_not_called()
-
- assert await hass.config_entries.async_unload(config_entry_2.entry_id)
- remove_service_mock.assert_not_called()
- assert await hass.config_entries.async_unload(config_entry.entry_id)
- assert remove_service_mock.call_count == 2
-
-
async def test_reconnect_client(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
@@ -144,7 +106,7 @@ async def test_reconnect_client_hub_unavailable(
config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_response=clients
)
- hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
hub.websocket.available = False
aioclient_mock.clear_requests()
@@ -293,7 +255,7 @@ async def test_remove_clients_hub_unavailable(
config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_all_response=clients
)
- hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+ hub = config_entry.runtime_data
hub.websocket.available = False
aioclient_mock.clear_requests()
--
GitLab