Use mysensors config entry async_on_unload (#139978)

* Use config entry on unload in mysensors

* Test mysensors config entry load and unload

* Fix docstring

homeassistant/components/mysensors/__init__.py

@@ -17,7 +17,6 @@ from .const import (
     DOMAIN,
     MYSENSORS_DISCOVERED_NODES,
     MYSENSORS_GATEWAYS,
-    MYSENSORS_ON_UNLOAD,
     PLATFORMS,
     DevId,
     DiscoveryInfo,
@@ -62,13 +61,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if not unload_ok:
         return False
 
-    key = MYSENSORS_ON_UNLOAD.format(entry.entry_id)
-    if key in hass.data[DOMAIN]:
-        for fnct in hass.data[DOMAIN][key]:
-            fnct()
-
-        hass.data[DOMAIN].pop(key)
-
     del hass.data[DOMAIN][MYSENSORS_GATEWAYS][entry.entry_id]
     hass.data[DOMAIN].pop(MYSENSORS_DISCOVERED_NODES.format(entry.entry_id), None)
 

homeassistant/components/mysensors/binary_sensor.py

@@ -20,7 +20,6 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 
 @dataclass(frozen=True)
@@ -86,9 +85,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.BINARY_SENSOR),

homeassistant/components/mysensors/climate.py

@@ -21,7 +21,6 @@ from homeassistant.util.unit_system import METRIC_SYSTEM
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 DICT_HA_TO_MYS = {
     HVACMode.AUTO: "AutoChangeOver",
@@ -57,9 +56,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.CLIMATE),

homeassistant/components/mysensors/const.py

@@ -34,7 +34,6 @@ CHILD_CALLBACK: str = "mysensors_child_callback_{}_{}_{}_{}"
 NODE_CALLBACK: str = "mysensors_node_callback_{}_{}"
 MYSENSORS_DISCOVERY: str = "mysensors_discovery_{}_{}"
 MYSENSORS_NODE_DISCOVERY: str = "mysensors_node_discovery"
-MYSENSORS_ON_UNLOAD: str = "mysensors_on_unload_{}"
 TYPE: Final = "type"
 UPDATE_DELAY: float = 0.1
 

homeassistant/components/mysensors/cover.py

@@ -15,7 +15,6 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 
 @unique
@@ -45,9 +44,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.COVER),

homeassistant/components/mysensors/device_tracker.py

@@ -12,7 +12,6 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 
 async def async_setup_entry(
@@ -33,9 +32,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.DEVICE_TRACKER),

homeassistant/components/mysensors/gateway.py

@@ -47,7 +47,6 @@ from .handler import HANDLERS
 from .helpers import (
     discover_mysensors_node,
     discover_mysensors_platform,
-    on_unload,
     validate_child,
     validate_node,
 )
@@ -293,9 +292,7 @@ async def _gw_start(
         """Stop the gateway."""
         await gw_stop(hass, entry, gateway)
 
-    on_unload(
-        hass,
-        entry.entry_id,
+    entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, stop_this_gw),
     )
 

homeassistant/components/mysensors/helpers.py

@@ -27,7 +27,6 @@ from .const import (
     MYSENSORS_DISCOVERED_NODES,
     MYSENSORS_DISCOVERY,
     MYSENSORS_NODE_DISCOVERY,
-    MYSENSORS_ON_UNLOAD,
     TYPE_TO_PLATFORMS,
     DevId,
     GatewayId,
@@ -41,18 +40,6 @@ SCHEMAS: Registry[
 ] = Registry()
 
 
-@callback
-def on_unload(hass: HomeAssistant, gateway_id: GatewayId, fnct: Callable) -> None:
-    """Register a callback to be called when entry is unloaded.
-
-    This function is used by platforms to cleanup after themselves.
-    """
-    key = MYSENSORS_ON_UNLOAD.format(gateway_id)
-    if key not in hass.data[DOMAIN]:
-        hass.data[DOMAIN][key] = []
-    hass.data[DOMAIN][key].append(fnct)
-
-
 @callback
 def discover_mysensors_platform(
     hass: HomeAssistant, gateway_id: GatewayId, platform: str, new_devices: list[DevId]

homeassistant/components/mysensors/light.py

@@ -21,7 +21,6 @@ from homeassistant.util.color import rgb_hex_to_rgb_list
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo, SensorType
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 
 async def async_setup_entry(
@@ -46,9 +45,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.LIGHT),

homeassistant/components/mysensors/remote.py

@@ -19,7 +19,6 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 
 async def async_setup_entry(
@@ -40,9 +39,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.REMOTE),

homeassistant/components/mysensors/sensor.py

@@ -50,7 +50,6 @@ from .const import (
     NodeDiscoveryInfo,
 )
 from .entity import MySensorNodeEntity, MySensorsChildEntity
-from .helpers import on_unload
 
 SENSORS: dict[str, SensorEntityDescription] = {
     "V_TEMP": SensorEntityDescription(
@@ -233,9 +232,7 @@ async def async_setup_entry(
         gateway: BaseAsyncGateway = hass.data[DOMAIN][MYSENSORS_GATEWAYS][gateway_id]
         async_add_entities([MyBatterySensor(gateway_id, gateway, node_id)])
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.SENSOR),
@@ -243,9 +240,7 @@ async def async_setup_entry(
         ),
     )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_NODE_DISCOVERY,

homeassistant/components/mysensors/switch.py

@@ -14,7 +14,6 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo, SensorType
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 
 async def async_setup_entry(
@@ -48,9 +47,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.SWITCH),

homeassistant/components/mysensors/text.py

@@ -12,7 +12,6 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from . import setup_mysensors_platform
 from .const import MYSENSORS_DISCOVERY, DiscoveryInfo
 from .entity import MySensorsChildEntity
-from .helpers import on_unload
 
 
 async def async_setup_entry(
@@ -33,9 +32,7 @@ async def async_setup_entry(
             async_add_entities=async_add_entities,
         )
 
-    on_unload(
-        hass,
-        config_entry.entry_id,
+    config_entry.async_on_unload(
         async_dispatcher_connect(
             hass,
             MYSENSORS_DISCOVERY.format(config_entry.entry_id, Platform.TEXT),

tests/components/mysensors/conftest.py

@@ -53,7 +53,7 @@ def gateway_nodes_fixture() -> dict[int, Sensor]:
 async def serial_transport_fixture(
     gateway_nodes: dict[int, Sensor],
     is_serial_port: MagicMock,
-) -> AsyncGenerator[dict[int, Sensor]]:
+) -> AsyncGenerator[MagicMock]:
     """Mock a serial transport."""
     with (
         patch(

tests/components/mysensors/test_init.py

@@ -2,10 +2,15 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
+from unittest.mock import MagicMock
+
 from mysensors import BaseSyncGateway
 from mysensors.sensor import Sensor
 
 from homeassistant.components.mysensors import DOMAIN
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
@@ -14,6 +19,50 @@ from tests.common import MockConfigEntry
 from tests.typing import WebSocketGenerator
 
 
+async def test_load_unload(
+    hass: HomeAssistant,
+    door_sensor: Sensor,
+    transport: MagicMock,
+    integration: MockConfigEntry,
+    receive_message: Callable[[str], None],
+) -> None:
+    """Test loading and unloading the MySensors config entry."""
+    config_entry = integration
+
+    assert config_entry.state == ConfigEntryState.LOADED
+
+    entity_id = "binary_sensor.door_sensor_1_1"
+    state = hass.states.get(entity_id)
+
+    assert state
+    assert state.state != STATE_UNAVAILABLE
+
+    receive_message("1;1;1;0;16;1\n")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+
+    assert state
+    assert state.state != STATE_UNAVAILABLE
+
+    assert await hass.config_entries.async_unload(config_entry.entry_id)
+
+    assert transport.return_value.disconnect.call_count == 1
+
+    state = hass.states.get(entity_id)
+
+    assert state
+    assert state.state == STATE_UNAVAILABLE
+
+    receive_message("1;1;1;0;16;1\n")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+
+    assert state
+    assert state.state == STATE_UNAVAILABLE
+
+
 async def test_remove_config_entry_device(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,