Use runtime_data in Axis integration (#116729)

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>

homeassistant/components/axis/__init__.py

@@ -13,8 +13,10 @@ from .hub import AxisHub, get_axis_api
 
 _LOGGER = logging.getLogger(__name__)
 
+AxisConfigEntry = ConfigEntry[AxisHub]
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, config_entry: AxisConfigEntry) -> bool:
     """Set up the Axis integration."""
     hass.data.setdefault(AXIS_DOMAIN, {})
 
@@ -25,8 +27,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     except AuthenticationRequired as err:
         raise ConfigEntryAuthFailed from err
 
-    hub = AxisHub(hass, config_entry, api)
-    hass.data[AXIS_DOMAIN][config_entry.entry_id] = hub
+    hub = config_entry.runtime_data = AxisHub(hass, config_entry, api)
     await hub.async_update_device_registry()
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
     hub.setup()
@@ -42,7 +43,6 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload Axis device config entry."""
-    hass.data[AXIS_DOMAIN].pop(config_entry.entry_id)
     return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
 
 

homeassistant/components/axis/binary_sensor.py

@@ -17,11 +17,11 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_call_later
 
+from . import AxisConfigEntry
 from .entity import AxisEventDescription, AxisEventEntity
 from .hub import AxisHub
 
@@ -177,11 +177,11 @@ ENTITY_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AxisConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a Axis binary sensor."""
-    AxisHub.get_hub(hass, config_entry).entity_loader.register_platform(
+    config_entry.runtime_data.entity_loader.register_platform(
         async_add_entities, AxisBinarySensor, ENTITY_DESCRIPTIONS
     )
 

homeassistant/components/axis/camera.py

@@ -4,12 +4,12 @@ from urllib.parse import urlencode
 
 from homeassistant.components.camera import CameraEntityFeature
 from homeassistant.components.mjpeg import MjpegCamera, filter_urllib3_logging
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import HTTP_DIGEST_AUTHENTICATION
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import AxisConfigEntry
 from .const import DEFAULT_STREAM_PROFILE, DEFAULT_VIDEO_SOURCE
 from .entity import AxisEntity
 from .hub import AxisHub
@@ -17,13 +17,13 @@ from .hub import AxisHub
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AxisConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Axis camera video stream."""
     filter_urllib3_logging()
 
-    hub = AxisHub.get_hub(hass, config_entry)
+    hub = config_entry.runtime_data
 
     if (
         not (prop := hub.api.vapix.params.property_handler.get("0"))

homeassistant/components/axis/config_flow.py

@@ -32,6 +32,7 @@ from homeassistant.core import callback
 from homeassistant.helpers.device_registry import format_mac
 from homeassistant.util.network import is_link_local
 
+from . import AxisConfigEntry
 from .const import (
     CONF_STREAM_PROFILE,
     CONF_VIDEO_SOURCE,
@@ -260,13 +261,14 @@ class AxisFlowHandler(ConfigFlow, domain=AXIS_DOMAIN):
 class AxisOptionsFlowHandler(OptionsFlowWithConfigEntry):
     """Handle Axis device options."""
 
+    config_entry: AxisConfigEntry
     hub: AxisHub
 
     async def async_step_init(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Manage the Axis device options."""
-        self.hub = AxisHub.get_hub(self.hass, self.config_entry)
+        self.hub = self.config_entry.runtime_data
         return await self.async_step_configure_stream()
 
     async def async_step_configure_stream(

homeassistant/components/axis/diagnostics.py

@@ -5,11 +5,10 @@ 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_MAC, CONF_PASSWORD, CONF_UNIQUE_ID, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from .hub import AxisHub
+from . import AxisConfigEntry
 
 REDACT_CONFIG = {CONF_MAC, CONF_PASSWORD, CONF_UNIQUE_ID, CONF_USERNAME}
 REDACT_BASIC_DEVICE_INFO = {"SerialNumber", "SocSerialNumber"}
@@ -17,10 +16,10 @@ REDACT_VAPIX_PARAMS = {"root.Network", "System.SerialNumber"}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: AxisConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    hub = AxisHub.get_hub(hass, config_entry)
+    hub = config_entry.runtime_data
     diag: dict[str, Any] = hub.additional_diagnostics.copy()
 
     diag["config"] = async_redact_data(config_entry.as_dict(), REDACT_CONFIG)

homeassistant/components/axis/hub/hub.py

@@ -2,11 +2,10 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import axis
 
-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 CONNECTION_NETWORK_MAC, format_mac
@@ -17,12 +16,15 @@ from .config import AxisConfig
 from .entity_loader import AxisEntityLoader
 from .event_source import AxisEventSource
 
+if TYPE_CHECKING:
+    from .. import AxisConfigEntry
+
 
 class AxisHub:
     """Manages a Axis device."""
 
     def __init__(
-        self, hass: HomeAssistant, config_entry: ConfigEntry, api: axis.AxisDevice
+        self, hass: HomeAssistant, config_entry: AxisConfigEntry, api: axis.AxisDevice
     ) -> None:
         """Initialize the device."""
         self.hass = hass
@@ -37,13 +39,6 @@ class AxisHub:
 
         self.additional_diagnostics: dict[str, Any] = {}
 
-    @callback
-    @staticmethod
-    def get_hub(hass: HomeAssistant, config_entry: ConfigEntry) -> AxisHub:
-        """Get Axis hub from config entry."""
-        hub: AxisHub = hass.data[AXIS_DOMAIN][config_entry.entry_id]
-        return hub
-
     @property
     def available(self) -> bool:
         """Connection state to the device."""
@@ -63,7 +58,7 @@ class AxisHub:
 
     @staticmethod
     async def async_new_address_callback(
-        hass: HomeAssistant, config_entry: ConfigEntry
+        hass: HomeAssistant, config_entry: AxisConfigEntry
     ) -> None:
         """Handle signals of device getting new address.
 
@@ -71,7 +66,7 @@ class AxisHub:
         This is a static method because a class method (bound method),
         cannot be used with weak references.
         """
-        hub = AxisHub.get_hub(hass, config_entry)
+        hub = config_entry.runtime_data
         hub.config = AxisConfig.from_config_entry(config_entry)
         hub.event_source.config_entry = config_entry
         hub.api.config.host = hub.config.host

homeassistant/components/axis/light.py

@@ -11,10 +11,10 @@ from homeassistant.components.light import (
     LightEntity,
     LightEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import AxisConfigEntry
 from .entity import TOPIC_TO_EVENT_TYPE, AxisEventDescription, AxisEventEntity
 from .hub import AxisHub
 
@@ -45,11 +45,11 @@ ENTITY_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AxisConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Axis light platform."""
-    AxisHub.get_hub(hass, config_entry).entity_loader.register_platform(
+    config_entry.runtime_data.entity_loader.register_platform(
         async_add_entities, AxisLight, ENTITY_DESCRIPTIONS
     )
 

homeassistant/components/axis/switch.py

@@ -10,11 +10,11 @@ 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.entity_platform import AddEntitiesCallback
 
+from . import AxisConfigEntry
 from .entity import AxisEventDescription, AxisEventEntity
 from .hub import AxisHub
 
@@ -38,11 +38,11 @@ ENTITY_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AxisConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Axis switch platform."""
-    AxisHub.get_hub(hass, config_entry).entity_loader.register_platform(
+    config_entry.runtime_data.entity_loader.register_platform(
         async_add_entities, AxisSwitch, ENTITY_DESCRIPTIONS
     )
 

tests/components/axis/test_hub.py

@@ -9,7 +9,6 @@ import pytest
 
 from homeassistant.components import axis, zeroconf
 from homeassistant.components.axis.const import DOMAIN as AXIS_DOMAIN
-from homeassistant.components.axis.hub import AxisHub
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
 from homeassistant.config_entries import SOURCE_ZEROCONF
 from homeassistant.const import (
@@ -52,7 +51,7 @@ async def test_device_setup(
     device_registry: dr.DeviceRegistry,
 ) -> None:
     """Successful setup."""
-    hub = AxisHub.get_hub(hass, setup_config_entry)
+    hub = setup_config_entry.runtime_data
 
     assert hub.api.vapix.firmware_version == "9.10.1"
     assert hub.api.vapix.product_number == "M1065-LW"
@@ -78,7 +77,7 @@ async def test_device_setup(
 @pytest.mark.parametrize("api_discovery_items", [API_DISCOVERY_BASIC_DEVICE_INFO])
 async def test_device_info(hass: HomeAssistant, setup_config_entry) -> None:
     """Verify other path of device information works."""
-    hub = AxisHub.get_hub(hass, setup_config_entry)
+    hub = setup_config_entry.runtime_data
 
     assert hub.api.vapix.firmware_version == "9.80.1"
     assert hub.api.vapix.product_number == "M1065-LW"
@@ -124,12 +123,9 @@ async def test_update_address(
     hass: HomeAssistant, setup_config_entry, mock_vapix_requests
 ) -> None:
     """Test update address works."""
-    hub = AxisHub.get_hub(hass, setup_config_entry)
+    hub = setup_config_entry.runtime_data
     assert hub.api.config.host == "1.2.3.4"
 
-    with patch(
-        "homeassistant.components.axis.async_setup_entry", return_value=True
-    ) as mock_setup_entry:
     mock_vapix_requests("2.3.4.5")
     await hass.config_entries.flow.async_init(
         AXIS_DOMAIN,
@@ -147,7 +143,6 @@ async def test_update_address(
     await hass.async_block_till_done()
 
     assert hub.api.config.host == "2.3.4.5"
-    assert len(mock_setup_entry.mock_calls) == 1
 
 
 async def test_device_unavailable(