Move runtime_data deletion after unload (#119224)

* Move runtime_data deletion after unload. Doing this before unload means we can't use, eg. the coordinator, during teardown. * Re-order config entry on unload * Add test --------- Co-authored-by: Paulus Schoutsen <[email protected]>
home-assistant · Jun 11, 2024 · 9bb9792 · 9bb9792
1 parent 3308f07
commit 9bb9792
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py
@@ -803,13 +803,13 @@ async def async_unload(
             assert isinstance(result, bool)
 
             # Only adjust state if we unloaded the component
-            if domain_is_integration:
-                if result:
-                    self._async_set_state(hass, ConfigEntryState.NOT_LOADED, None)
-                    if hasattr(self, "runtime_data"):
-                        object.__delattr__(self, "runtime_data")
-
+            if domain_is_integration and result:
                 await self._async_process_on_unload(hass)
+                if hasattr(self, "runtime_data"):
+                    object.__delattr__(self, "runtime_data")
+
+                self._async_set_state(hass, ConfigEntryState.NOT_LOADED, None)
+
         except Exception as exc:
             _LOGGER.exception(
                 "Error unloading entry %s for %s", self.title, integration.domain

diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py
@@ -1726,16 +1726,23 @@ async def test_entry_unload_succeed(
     hass: HomeAssistant, manager: config_entries.ConfigEntries
 ) -> None:
     """Test that we can unload an entry."""
+    unloads_called = []
+
+    async def verify_runtime_data(*args):
+        """Verify runtime data."""
+        assert entry.runtime_data == 2
+        unloads_called.append(args)
+        return True
+
     entry = MockConfigEntry(domain="comp", state=config_entries.ConfigEntryState.LOADED)
     entry.add_to_hass(hass)
+    entry.async_on_unload(verify_runtime_data)
     entry.runtime_data = 2
 
-    async_unload_entry = AsyncMock(return_value=True)
-
-    mock_integration(hass, MockModule("comp", async_unload_entry=async_unload_entry))
+    mock_integration(hass, MockModule("comp", async_unload_entry=verify_runtime_data))
 
     assert await manager.async_unload(entry.entry_id)
-    assert len(async_unload_entry.mock_calls) == 1
+    assert len(unloads_called) == 2
     assert entry.state is config_entries.ConfigEntryState.NOT_LOADED
     assert not hasattr(entry, "runtime_data")