diff --git a/homeassistant/components/agent_dvr/__init__.py b/homeassistant/components/agent_dvr/__init__.py index 6dc83d3766d..2cb32b6c80e 100644 --- a/homeassistant/components/agent_dvr/__init__.py +++ b/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) diff --git a/homeassistant/components/agent_dvr/alarm_control_panel.py b/homeassistant/components/agent_dvr/alarm_control_panel.py index 8dae49aa0ea..e703bcad6ae 100644 --- a/homeassistant/components/agent_dvr/alarm_control_panel.py +++ b/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): diff --git a/homeassistant/components/agent_dvr/camera.py b/homeassistant/components/agent_dvr/camera.py index 88ffd8bcc39..4438bf72a1a 100644 --- a/homeassistant/components/agent_dvr/camera.py +++ b/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 diff --git a/homeassistant/components/agent_dvr/const.py b/homeassistant/components/agent_dvr/const.py index cd0284ca87c..8557f0595ed 100644 --- a/homeassistant/components/agent_dvr/const.py +++ b/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" diff --git a/tests/components/agent_dvr/test_init.py b/tests/components/agent_dvr/test_init.py index 7f546a190a7..5e263c548c8 100644 --- a/tests/components/agent_dvr/test_init.py +++ b/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: