From efd7ddeb89f643241ffe78680bfb00b063de4667 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 14 Feb 2025 13:06:07 +0100
Subject: [PATCH] Improve tests of removing and unloading config entries
 (#138432)

* Improve tests of removing and unloading config entries

* Fix unnecessary coroutine
---
 tests/test_config_entries.py | 54 +++++++++++++++++++++++++++---------
 1 file changed, 41 insertions(+), 13 deletions(-)

diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py
index cf022c42e94..bf2280790fa 100644
--- a/tests/test_config_entries.py
+++ b/tests/test_config_entries.py
@@ -462,7 +462,15 @@ async def test_remove_entry(
         assert result
         return result
 
-    mock_remove_entry = AsyncMock(return_value=None)
+    remove_entry_calls = []
+
+    async def mock_remove_entry(
+        hass: HomeAssistant, entry: config_entries.ConfigEntry
+    ) -> None:
+        """Mock removing an entry."""
+        # Check that the entry is not yet removed from config entries
+        assert hass.config_entries.async_get_entry(entry.entry_id)
+        remove_entry_calls.append(None)
 
     entity = MockEntity(unique_id="1234", name="Test Entity")
 
@@ -522,7 +530,7 @@ async def test_remove_entry(
     assert result == {"require_restart": False}
 
     # Check the remove callback was invoked.
-    assert mock_remove_entry.call_count == 1
+    assert len(remove_entry_calls) == 1
 
     # Check that config entry was removed.
     assert manager.async_entry_ids() == ["test1", "test3"]
@@ -2611,29 +2619,49 @@ async def test_entry_setup_invalid_state(
     assert entry.state is state
 
 
-async def test_entry_unload_succeed(
-    hass: HomeAssistant, manager: config_entries.ConfigEntries
+@pytest.mark.parametrize(
+    ("unload_result", "expected_result", "expected_state", "has_runtime_data"),
+    [
+        (True, True, config_entries.ConfigEntryState.NOT_LOADED, False),
+        (False, False, config_entries.ConfigEntryState.LOADED, True),
+    ],
+)
+async def test_entry_unload(
+    hass: HomeAssistant,
+    manager: config_entries.ConfigEntries,
+    unload_result: bool,
+    expected_result: bool,
+    expected_state: config_entries.ConfigEntryState,
+    has_runtime_data: bool,
 ) -> None:
     """Test that we can unload an entry."""
-    unloads_called = []
+    unload_entry_calls = []
 
-    async def verify_runtime_data(*args):
+    @callback
+    def verify_runtime_data() -> None:
         """Verify runtime data."""
         assert entry.runtime_data == 2
-        unloads_called.append(args)
-        return True
+
+    async def async_unload_entry(
+        hass: HomeAssistant, entry: config_entries.ConfigEntry
+    ) -> bool:
+        """Mock unload entry."""
+        unload_entry_calls.append(None)
+        verify_runtime_data()
+        assert entry.state is config_entries.ConfigEntryState.LOADED
+        return unload_result
 
     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
 
-    mock_integration(hass, MockModule("comp", async_unload_entry=verify_runtime_data))
+    mock_integration(hass, MockModule("comp", async_unload_entry=async_unload_entry))
 
-    assert await manager.async_unload(entry.entry_id)
-    assert len(unloads_called) == 2
-    assert entry.state is config_entries.ConfigEntryState.NOT_LOADED
-    assert not hasattr(entry, "runtime_data")
+    assert await manager.async_unload(entry.entry_id) == expected_result
+    assert len(unload_entry_calls) == 1
+    assert entry.state is expected_state
+    assert hasattr(entry, "runtime_data") == has_runtime_data
 
 
 @pytest.mark.parametrize(