Use ConfigEntry runtime_data in agent_dvr (#117412)

homeassistant/components/agent_dvr/__init__.py

@@ -10,18 +10,20 @@ from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import CONNECTION, DOMAIN as AGENT_DOMAIN, SERVER_URL
+from .const import DOMAIN as AGENT_DOMAIN, SERVER_URL
 
 ATTRIBUTION = "ispyconnect.com"
 DEFAULT_BRAND = "Agent DVR by ispyconnect.com"
 
 PLATFORMS = [Platform.ALARM_CONTROL_PANEL, Platform.CAMERA]
 
+AgentDVRConfigEntry = ConfigEntry[Agent]
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(
+    hass: HomeAssistant, config_entry: AgentDVRConfigEntry
+) -> bool:
     """Set up the Agent component."""
-    hass.data.setdefault(AGENT_DOMAIN, {})
-
     server_origin = config_entry.data[SERVER_URL]
 
     agent_client = Agent(server_origin, async_get_clientsession(hass))
@@ -34,9 +36,11 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     if not agent_client.is_available:
         raise ConfigEntryNotReady
 
+    config_entry.async_on_unload(agent_client.close)
+
     await agent_client.get_devices()
 
-    hass.data[AGENT_DOMAIN][config_entry.entry_id] = {CONNECTION: agent_client}
+    config_entry.runtime_data = agent_client
 
     device_registry = dr.async_get(hass)
 
@@ -54,15 +58,8 @@ 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: AgentDVRConfigEntry
+) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(
-        config_entry, PLATFORMS
-    )
-
-    await hass.data[AGENT_DOMAIN][config_entry.entry_id][CONNECTION].close()
-
-    if unload_ok:
-        hass.data[AGENT_DOMAIN].pop(config_entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)

homeassistant/components/agent_dvr/alarm_control_panel.py

@@ -6,7 +6,6 @@ from homeassistant.components.alarm_control_panel import (
     AlarmControlPanelEntity,
     AlarmControlPanelEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     STATE_ALARM_ARMED_AWAY,
     STATE_ALARM_ARMED_HOME,
@@ -17,7 +16,8 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import CONNECTION, DOMAIN as AGENT_DOMAIN
+from . import AgentDVRConfigEntry
+from .const import DOMAIN as AGENT_DOMAIN
 
 CONF_HOME_MODE_NAME = "home"
 CONF_AWAY_MODE_NAME = "away"
@@ -28,13 +28,11 @@ CONST_ALARM_CONTROL_PANEL_NAME = "Alarm Panel"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AgentDVRConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Agent DVR Alarm Control Panels."""
-    async_add_entities(
-        [AgentBaseStation(hass.data[AGENT_DOMAIN][config_entry.entry_id][CONNECTION])]
-    )
+    async_add_entities([AgentBaseStation(config_entry.runtime_data)])
 
 
 class AgentBaseStation(AlarmControlPanelEntity):

homeassistant/components/agent_dvr/camera.py

@@ -7,7 +7,6 @@ from agent import AgentError
 
 from homeassistant.components.camera import CameraEntityFeature
 from homeassistant.components.mjpeg import MjpegCamera, filter_urllib3_logging
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import (
@@ -15,12 +14,8 @@ from homeassistant.helpers.entity_platform import (
     async_get_current_platform,
 )
 
-from .const import (
-    ATTRIBUTION,
-    CAMERA_SCAN_INTERVAL_SECS,
-    CONNECTION,
-    DOMAIN as AGENT_DOMAIN,
-)
+from . import AgentDVRConfigEntry
+from .const import ATTRIBUTION, CAMERA_SCAN_INTERVAL_SECS, DOMAIN as AGENT_DOMAIN
 
 SCAN_INTERVAL = timedelta(seconds=CAMERA_SCAN_INTERVAL_SECS)
 
@@ -43,14 +38,14 @@ CAMERA_SERVICES = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AgentDVRConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Agent cameras."""
     filter_urllib3_logging()
     cameras = []
 
-    server = hass.data[AGENT_DOMAIN][config_entry.entry_id][CONNECTION]
+    server = config_entry.runtime_data
     if not server.devices:
         _LOGGER.warning("Could not fetch cameras from Agent server")
         return

homeassistant/components/agent_dvr/const.py

@@ -9,4 +9,3 @@ SERVICE_UPDATE = "update"
 SIGNAL_UPDATE_AGENT = "agent_update"
 ATTRIBUTION = "Data provided by ispyconnect.com"
 SERVER_URL = "server_url"
-CONNECTION = "connection"

tests/components/agent_dvr/test_init.py

@@ -39,7 +39,6 @@ async def test_setup_config_and_unload(
     await hass.async_block_till_done()
 
     assert entry.state is ConfigEntryState.NOT_LOADED
-    assert not hass.data.get(DOMAIN)
 
 
 async def test_async_setup_entry_not_ready(hass: HomeAssistant) -> None: