Fix shadow observer lifecycle, learning and prediction correctness

custom_components/heating_mpc/__init__.py

@@ -5,8 +5,10 @@ from homeassistant.const import Platform
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.storage import Store
+from homeassistant.helpers import entity_registry as er
 
-from .const import DOMAIN
+from .const import DOMAIN, CONF_OUTSIDE, WEATHER_SOURCES
+from .weather_sources import sources_for
 from .coordinator import HeatingMPCCoordinator
 
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.BINARY_SENSOR]
@@ -17,7 +19,20 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     await coordinator.async_restore_model()
     await coordinator.async_config_entry_first_refresh()
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
-    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    try:
+        await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    except Exception:
+        hass.data[DOMAIN].pop(entry.entry_id, None)
+        raise
+    # Remove only obsolete room forecasts belonging to this entry/platform.
+    registry = er.async_get(hass)
+    expected = {f"{entry.entry_id}_{room}_predicted_60min" for room in coordinator.rooms}
+    for entity in er.async_entries_for_config_entry(registry, entry.entry_id):
+        if (entity.platform == DOMAIN and
+                entity.unique_id.startswith(f"{entry.entry_id}_sensor.") and
+                entity.unique_id.endswith("_predicted_60min") and
+                entity.unique_id not in expected):
+            registry.async_remove(entity.entity_id)
     entry.async_on_unload(entry.add_update_listener(async_reload_entry))
     return True
 
@@ -37,3 +52,21 @@ async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Remove the learned model when this integration instance is deleted."""
     store = Store(hass, 1, f"{DOMAIN}.{entry.entry_id}")
     await store.async_remove()
+
+
+async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Migrate legacy single weather source without changing user selections."""
+    if entry.version > 2:
+        return False
+    if entry.version == 1:
+        def migrate(config):
+            result = dict(config)
+            if CONF_OUTSIDE in result:
+                result[WEATHER_SOURCES["temperature"]] = sources_for(result, "temperature")
+                result.pop(CONF_OUTSIDE)
+            return result
+        hass.config_entries.async_update_entry(
+            entry, data=migrate(entry.data), options=migrate(entry.options),
+            version=2, minor_version=1,
+        )
+    return True