Store runtime data inside the config entry in Speedtest.net (#116802)

homeassistant/components/speedtestdotnet/__init__.py

@@ -12,13 +12,16 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.start import async_at_started
 
-from .const import DOMAIN
 from .coordinator import SpeedTestDataCoordinator
 
 PLATFORMS = [Platform.SENSOR]
 
+SpeedTestConfigEntry = ConfigEntry[SpeedTestDataCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(
+    hass: HomeAssistant, config_entry: SpeedTestConfigEntry
+) -> bool:
     """Set up the Speedtest.net component."""
     try:
         api = await hass.async_add_executor_job(
@@ -28,7 +31,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     except speedtest.SpeedtestException as err:
         raise ConfigEntryNotReady from err
 
-    hass.data[DOMAIN] = coordinator
+    config_entry.runtime_data = coordinator
 
     async def _async_finish_startup(hass: HomeAssistant) -> None:
         """Run this only when HA has finished its startup."""
@@ -45,11 +48,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload SpeedTest Entry from config_entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(
-        config_entry, PLATFORMS
-    ):
-        hass.data.pop(DOMAIN)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
 
 
 async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> None:

homeassistant/components/speedtestdotnet/config_flow.py

@@ -6,14 +6,10 @@ from typing import Any
 
 import voluptuous as vol
 
-from homeassistant.config_entries import (
-    ConfigEntry,
-    ConfigFlow,
-    ConfigFlowResult,
-    OptionsFlow,
-)
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult, OptionsFlow
 from homeassistant.core import callback
 
+from . import SpeedTestConfigEntry
 from .const import (
     CONF_SERVER_ID,
     CONF_SERVER_NAME,
@@ -31,7 +27,7 @@ class SpeedTestFlowHandler(ConfigFlow, domain=DOMAIN):
     @staticmethod
     @callback
     def async_get_options_flow(
-        config_entry: ConfigEntry,
+        config_entry: SpeedTestConfigEntry,
     ) -> SpeedTestOptionsFlowHandler:
         """Get the options flow for this handler."""
         return SpeedTestOptionsFlowHandler(config_entry)
@@ -52,7 +48,7 @@ class SpeedTestFlowHandler(ConfigFlow, domain=DOMAIN):
 class SpeedTestOptionsFlowHandler(OptionsFlow):
     """Handle SpeedTest options."""
 
-    def __init__(self, config_entry: ConfigEntry) -> None:
+    def __init__(self, config_entry: SpeedTestConfigEntry) -> None:
         """Initialize options flow."""
         self.config_entry = config_entry
         self._servers: dict = {}
@@ -73,7 +69,7 @@ class SpeedTestOptionsFlowHandler(OptionsFlow):
 
             return self.async_create_entry(title="", data=user_input)
 
-        self._servers = self.hass.data[DOMAIN].servers
+        self._servers = self.config_entry.runtime_data.servers
 
         options = {
             vol.Optional(

homeassistant/components/speedtestdotnet/sensor.py

@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfDataRate, UnitOfTime
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
@@ -20,6 +19,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import SpeedTestConfigEntry
 from .const import (
     ATTR_BYTES_RECEIVED,
     ATTR_BYTES_SENT,
@@ -69,11 +69,11 @@ SENSOR_TYPES: tuple[SpeedtestSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: SpeedTestConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Speedtestdotnet sensors."""
-    speedtest_coordinator = hass.data[DOMAIN]
+    speedtest_coordinator = config_entry.runtime_data
     async_add_entities(
         SpeedtestSensor(speedtest_coordinator, description)
         for description in SENSOR_TYPES

tests/components/speedtestdotnet/test_config_flow.py

@@ -3,7 +3,6 @@
 from unittest.mock import MagicMock
 
 from homeassistant import config_entries
-from homeassistant.components import speedtestdotnet
 from homeassistant.components.speedtestdotnet.const import (
     CONF_SERVER_ID,
     CONF_SERVER_NAME,
@@ -18,7 +17,7 @@ from tests.common import MockConfigEntry
 async def test_flow_works(hass: HomeAssistant) -> None:
     """Test user config."""
     result = await hass.config_entries.flow.async_init(
-        speedtestdotnet.DOMAIN, context={"source": config_entries.SOURCE_USER}
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "user"
@@ -84,7 +83,7 @@ async def test_integration_already_configured(hass: HomeAssistant) -> None:
     )
     entry.add_to_hass(hass)
     result = await hass.config_entries.flow.async_init(
-        speedtestdotnet.DOMAIN, context={"source": config_entries.SOURCE_USER}
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
     assert result["type"] is FlowResultType.ABORT
     assert result["reason"] == "single_instance_allowed"

tests/components/speedtestdotnet/test_init.py

@@ -10,6 +10,9 @@ from homeassistant.components.speedtestdotnet.const import (
     CONF_SERVER_NAME,
     DOMAIN,
 )
+from homeassistant.components.speedtestdotnet.coordinator import (
+    SpeedTestDataCoordinator,
+)
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
@@ -47,13 +50,12 @@ async def test_entry_lifecycle(hass: HomeAssistant, mock_api: MagicMock) -> None
     await hass.async_block_till_done()
 
     assert entry.state is ConfigEntryState.LOADED
-    assert hass.data[DOMAIN]
+    assert isinstance(entry.runtime_data, SpeedTestDataCoordinator)
 
     assert await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()
 
     assert entry.state is ConfigEntryState.NOT_LOADED
-    assert DOMAIN not in hass.data
 
 
 async def test_server_not_found(hass: HomeAssistant, mock_api: MagicMock) -> None:
@@ -67,7 +69,9 @@ async def test_server_not_found(hass: HomeAssistant, mock_api: MagicMock) -> Non
 
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
-    assert hass.data[DOMAIN]
+
+    assert entry.state is ConfigEntryState.LOADED
+    assert isinstance(entry.runtime_data, SpeedTestDataCoordinator)
 
     mock_api.return_value.get_servers.side_effect = speedtest.NoMatchedServers
     async_fire_time_changed(
@@ -90,14 +94,16 @@ async def test_get_best_server_error(hass: HomeAssistant, mock_api: MagicMock) -
 
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()
-    assert hass.data[DOMAIN]
+
+    assert entry.state is ConfigEntryState.LOADED
+    assert isinstance(entry.runtime_data, SpeedTestDataCoordinator)
 
     mock_api.return_value.get_best_server.side_effect = (
         speedtest.SpeedtestBestServerFailure(
             "Unable to connect to servers to test latency."
         )
     )
-    await hass.data[DOMAIN].async_refresh()
+    await entry.runtime_data.async_refresh()
     await hass.async_block_till_done()
     state = hass.states.get("sensor.speedtest_ping")
     assert state is not None

tests/components/speedtestdotnet/test_sensor.py

@@ -3,7 +3,7 @@
 from unittest.mock import MagicMock
 
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
-from homeassistant.components.speedtestdotnet import DOMAIN
+from homeassistant.components.speedtestdotnet.const import DOMAIN
 from homeassistant.core import HomeAssistant
 
 from . import MOCK_RESULTS, MOCK_SERVERS, MOCK_STATES