Store runtime data inside the config entry in IPP (#116765)

* store runtime data inside the config entry

* improve tests

homeassistant/components/ipp/__init__.py

@@ -12,13 +12,15 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 
-from .const import CONF_BASE_PATH, DOMAIN
+from .const import CONF_BASE_PATH
 from .coordinator import IPPDataUpdateCoordinator
 
 PLATFORMS = [Platform.SENSOR]
 
+IPPConfigEntry = ConfigEntry[IPPDataUpdateCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: IPPConfigEntry) -> bool:
     """Set up IPP from a config entry."""
     # config flow sets this to either UUID, serial number or None
     if (device_id := entry.unique_id) is None:
@@ -35,7 +37,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -44,6 +46,4 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok

homeassistant/components/ipp/diagnostics.py

@@ -4,18 +4,16 @@ from __future__ import annotations
 
 from typing import Any
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
+from . import IPPConfigEntry
-from .coordinator import IPPDataUpdateCoordinator
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: IPPConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: IPPDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
 
     return {
         "entry": {

homeassistant/components/ipp/sensor.py

@@ -15,13 +15,13 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_LOCATION, PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.util.dt import utcnow
 
+from . import IPPConfigEntry
 from .const import (
     ATTR_COMMAND_SET,
     ATTR_INFO,
@@ -32,9 +32,7 @@ from .const import (
     ATTR_STATE_MESSAGE,
     ATTR_STATE_REASON,
     ATTR_URI_SUPPORTED,
-    DOMAIN,
 )
-from .coordinator import IPPDataUpdateCoordinator
 from .entity import IPPEntity
 
 
@@ -89,11 +87,11 @@ PRINTER_SENSORS: tuple[IPPSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: IPPConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up IPP sensor based on a config entry."""
-    coordinator: IPPDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     sensors: list[SensorEntity] = [
         IPPSensor(
             coordinator,

tests/components/ipp/test_init.py

@@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 from pyipp import IPPConnectionError
 
-from homeassistant.components.ipp.const import DOMAIN
+from homeassistant.components.ipp.coordinator import IPPDataUpdateCoordinator
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -37,10 +37,9 @@ async def test_load_unload_config_entry(
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert mock_config_entry.entry_id in hass.data[DOMAIN]
     assert mock_config_entry.state is ConfigEntryState.LOADED
+    assert isinstance(mock_config_entry.runtime_data, IPPDataUpdateCoordinator)
 
     await hass.config_entries.async_unload(mock_config_entry.entry_id)
     await hass.async_block_till_done()
-    assert mock_config_entry.entry_id not in hass.data[DOMAIN]
     assert mock_config_entry.state is ConfigEntryState.NOT_LOADED