home-assistant · frenck · Jun 14, 2021 · Mar 15, 2021 · Mar 15, 2021 · Mar 15, 2021
@@ -51,6 +51,30 @@ async def async_setup_entry(
     )
 
 
+def get_platforms(config_entry):
+    """Return the platforms belonging to a config_entry."""
+    model = config_entry.data[CONF_MODEL]
+    flow_type = config_entry.data[CONF_FLOW_TYPE]
+
+    if flow_type == CONF_GATEWAY:
+        return GATEWAY_PLATFORMS
+    if flow_type == CONF_DEVICE:
+        if model in MODELS_SWITCH:
+            return SWITCH_PLATFORMS
+        if model in MODELS_FAN:
+            return FAN_PLATFORMS
+        if model in MODELS_LIGHT:
+            return LIGHT_PLATFORMS
+        for vacuum_model in MODELS_VACUUM:
+            if model.startswith(vacuum_model):
+                return VACUUM_PLATFORMS
+        for air_monitor_model in MODELS_AIR_MONITOR:
+            if model.startswith(air_monitor_model):
+                return AIR_MONITOR_PLATFORMS
+
+    return []
+
+
 async def async_setup_gateway_entry(
     hass: core.HomeAssistant, entry: config_entries.ConfigEntry
 ):
@@ -64,8 +88,10 @@ async def async_setup_gateway_entry(
     if entry.unique_id.endswith("-gateway"):
         hass.config_entries.async_update_entry(entry, unique_id=entry.data["mac"])
 
+    entry.async_on_unload(entry.add_update_listener(update_listener))
+
     # Connect to gateway
-    gateway = ConnectXiaomiGateway(hass)
+    gateway = ConnectXiaomiGateway(hass, entry)
     if not await gateway.async_connect_gateway(host, token):
         return False
     gateway_info = gateway.gateway_info
@@ -128,29 +154,36 @@ async def async_setup_device_entry(
     hass: core.HomeAssistant, entry: config_entries.ConfigEntry
 ):
     """Set up the Xiaomi Miio device component from a config entry."""
-    model = entry.data[CONF_MODEL]
-
-    # Identify platforms to setup
-    platforms = []
-    if model in MODELS_SWITCH:
-        platforms = SWITCH_PLATFORMS
-    elif model in MODELS_FAN:
-        platforms = FAN_PLATFORMS
-    elif model in MODELS_LIGHT:
-        platforms = LIGHT_PLATFORMS
-    for vacuum_model in MODELS_VACUUM:
-        if model.startswith(vacuum_model):
-            platforms = VACUUM_PLATFORMS
-    for air_monitor_model in MODELS_AIR_MONITOR:
-        if model.startswith(air_monitor_model):
-            platforms = AIR_MONITOR_PLATFORMS
+    platforms = get_platforms(entry)
 
     if not platforms:
         return False
 
-    for platform in platforms:
-        hass.async_create_task(
-            hass.config_entries.async_forward_entry_setup(entry, platform)
-        )
+    entry.async_on_unload(entry.add_update_listener(update_listener))
+
+    hass.config_entries.async_setup_platforms(entry, platforms)
 
     return True
+
+
+async def async_unload_entry(
+    hass: core.HomeAssistant, config_entry: config_entries.ConfigEntry
+):
+    """Unload a config entry."""
+    platforms = get_platforms(config_entry)
+
+    unload_ok = await hass.config_entries.async_unload_platforms(
+        config_entry, platforms
+    )
+
+    if unload_ok:
+        hass.data[DOMAIN].pop(config_entry.entry_id)
+
+    return unload_ok
+
+
+async def update_listener(
+    hass: core.HomeAssistant, config_entry: config_entries.ConfigEntry
+):
+    """Handle options update."""
+    await hass.config_entries.async_reload(config_entry.entry_id)