diff --git a/homeassistant/components/speedtestdotnet/__init__.py b/homeassistant/components/speedtestdotnet/__init__.py index 3c15f2fb820..19525ad9bfa 100644 --- a/homeassistant/components/speedtestdotnet/__init__.py +++ b/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: diff --git a/homeassistant/components/speedtestdotnet/config_flow.py b/homeassistant/components/speedtestdotnet/config_flow.py index 2ef2a70d745..dc64448bbef 100644 --- a/homeassistant/components/speedtestdotnet/config_flow.py +++ b/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( diff --git a/homeassistant/components/speedtestdotnet/sensor.py b/homeassistant/components/speedtestdotnet/sensor.py index 5bf1a6bea91..10da1dc93af 100644 --- a/homeassistant/components/speedtestdotnet/sensor.py +++ b/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 diff --git a/tests/components/speedtestdotnet/test_config_flow.py b/tests/components/speedtestdotnet/test_config_flow.py index f509c91ad20..883f60aaf0a 100644 --- a/tests/components/speedtestdotnet/test_config_flow.py +++ b/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" diff --git a/tests/components/speedtestdotnet/test_init.py b/tests/components/speedtestdotnet/test_init.py index 446ed527df4..2e20aaa259c 100644 --- a/tests/components/speedtestdotnet/test_init.py +++ b/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 diff --git a/tests/components/speedtestdotnet/test_sensor.py b/tests/components/speedtestdotnet/test_sensor.py index e529d46b537..a14a482b66f 100644 --- a/tests/components/speedtestdotnet/test_sensor.py +++ b/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