From 021878e942e1e49785533b9b9ada8e67db90d598 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 9 Sep 2024 00:01:45 +0200
Subject: [PATCH] Migrate ambient_network to use runtime_data (#125535)

---
 .../components/ambient_network/__init__.py     | 18 ++++++++++--------
 .../components/ambient_network/sensor.py       |  7 +++----
 tests/components/ambient_network/conftest.py   |  4 ++--
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/homeassistant/components/ambient_network/__init__.py b/homeassistant/components/ambient_network/__init__.py
index b286fb7fbc9..e9443a676b5 100644
--- a/homeassistant/components/ambient_network/__init__.py
+++ b/homeassistant/components/ambient_network/__init__.py
@@ -8,28 +8,30 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
 from .coordinator import AmbientNetworkDataUpdateCoordinator
 
+type AmbientNetworkConfigEntry = ConfigEntry[AmbientNetworkDataUpdateCoordinator]
+
 PLATFORMS: list[Platform] = [Platform.SENSOR]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, entry: AmbientNetworkConfigEntry
+) -> bool:
     """Set up the Ambient Weather Network from a config entry."""
 
     api = OpenAPI()
     coordinator = AmbientNetworkDataUpdateCoordinator(hass, api)
     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)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, entry: AmbientNetworkConfigEntry
+) -> bool:
     """Unload a config entry."""
 
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
diff --git a/homeassistant/components/ambient_network/sensor.py b/homeassistant/components/ambient_network/sensor.py
index 132fc7dbd0d..336745f88ff 100644
--- a/homeassistant/components/ambient_network/sensor.py
+++ b/homeassistant/components/ambient_network/sensor.py
@@ -10,7 +10,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
     CONCENTRATION_PARTS_PER_MILLION,
@@ -29,7 +28,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from .const import DOMAIN
+from . import AmbientNetworkConfigEntry
 from .coordinator import AmbientNetworkDataUpdateCoordinator
 from .entity import AmbientNetworkEntity
 
@@ -271,12 +270,12 @@ SENSOR_DESCRIPTIONS = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: AmbientNetworkConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Ambient Network sensor entities."""
 
-    coordinator: AmbientNetworkDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     if coordinator.config_entry is not None:
         async_add_entities(
             AmbientNetworkSensor(
diff --git a/tests/components/ambient_network/conftest.py b/tests/components/ambient_network/conftest.py
index 9fc001252a0..e728d46aaf6 100644
--- a/tests/components/ambient_network/conftest.py
+++ b/tests/components/ambient_network/conftest.py
@@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, Mock, patch
 from aioambient import OpenAPI
 import pytest
 
-from homeassistant.components import ambient_network
+from homeassistant.components.ambient_network.const import DOMAIN
 from homeassistant.core import HomeAssistant
 
 from tests.common import (
@@ -69,7 +69,7 @@ async def mock_aioambient(open_api: OpenAPI):
 def config_entry_fixture(request: pytest.FixtureRequest) -> MockConfigEntry:
     """Mock config entry."""
     return MockConfigEntry(
-        domain=ambient_network.DOMAIN,
+        domain=DOMAIN,
         title=f"Station {request.param[0]}",
         data={"mac": request.param},
     )