diff --git a/tests/test_integration.py b/tests/test_integration.py
index aa7c378..968ba76 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -72,7 +72,7 @@ async def test_setup_entry_success(hass: HomeAssistant, mock_config_entry, mock_
     """Test successful setup of config entry."""
     with patch("custom_components.adguard_hub.AdGuardHomeAPI", return_value=mock_api), \
          patch("custom_components.adguard_hub.async_get_clientsession"), \
-         patch.object(hass.config_entries, "async_forward_entry_setups", return_value=True):
+         patch.object(hass.config_entries, "async_forward_entry_setups", AsyncMock(return_value=True)):
 
         result = await async_setup_entry(hass, mock_config_entry)
 
@@ -107,11 +107,12 @@ async def test_unload_entry(hass: HomeAssistant, mock_config_entry):
         }
     }
 
-    with patch.object(hass.config_entries, "async_unload_platforms", return_value=True):
+    with patch.object(hass.config_entries, "async_unload_platforms", AsyncMock(return_value=True)):
         result = await async_unload_entry(hass, mock_config_entry)
 
         assert result is True
-        assert mock_config_entry.entry_id not in hass.data[DOMAIN]
+        # FIXED: Handle case where DOMAIN key is removed when last entry is unloaded
+        assert DOMAIN not in hass.data or mock_config_entry.entry_id not in hass.data[DOMAIN]
 
 
 @pytest.mark.asyncio