Clean up patching of setup platforms for UniFi Protect integration (#63156)

tests/components/unifiprotect/conftest.py

@@ -13,7 +13,8 @@ import pytest
 from pyunifiprotect.data import Camera, Light, Version, WSSubscriptionMessage
 
 from homeassistant.components.unifiprotect.const import DOMAIN, MIN_REQUIRED_PROTECT_V
-from homeassistant.core import HomeAssistant
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant, split_entity_id
 from homeassistant.helpers import entity_registry as er
 import homeassistant.util.dt as dt_util
 
@@ -160,3 +161,18 @@ async def enable_entity(
     await hass.async_block_till_done()
 
     return updated_entity
+
+
+def assert_entity_counts(
+    hass: HomeAssistant, platform: Platform, total: int, enabled: int
+) -> None:
+    """Assert entity counts for a given platform."""
+
+    entity_registry = er.async_get(hass)
+
+    entities = [
+        e for e in entity_registry.entities if split_entity_id(e)[0] == platform.value
+    ]
+
+    assert len(entities) == total
+    assert len(hass.states.async_all(platform.value)) == enabled

tests/components/unifiprotect/test_button.py

@@ -2,7 +2,7 @@
 # pylint: disable=protected-access
 from __future__ import annotations
 
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock
 
 import pytest
 from pyunifiprotect.data import Camera
@@ -12,7 +12,7 @@ from homeassistant.const import ATTR_ATTRIBUTION, ATTR_ENTITY_ID, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from .conftest import MockEntityFixture, enable_entity
+from .conftest import MockEntityFixture, assert_entity_counts, enable_entity
 
 
 @pytest.fixture(name="camera")
@@ -32,14 +32,10 @@ async def camera_fixture(
         camera_obj.id: camera_obj,
     }
 
-    with patch("homeassistant.components.unifiprotect.PLATFORMS", [Platform.BUTTON]):
+    await hass.config_entries.async_setup(mock_entry.entry.entry_id)
-        await hass.config_entries.async_setup(mock_entry.entry.entry_id)
+    await hass.async_block_till_done()
-        await hass.async_block_till_done()
 
-    entity_registry = er.async_get(hass)
+    assert_entity_counts(hass, Platform.BUTTON, 1, 0)
-
-    assert len(hass.states.async_all()) == 0
-    assert len(entity_registry.entities) == 1
 
     yield (camera_obj, "button.test_camera_reboot_device")
 

tests/components/unifiprotect/test_camera.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from copy import copy
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 from pyunifiprotect.data import Camera as ProtectCamera
@@ -36,7 +36,12 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
 
-from .conftest import MockEntityFixture, enable_entity, time_changed
+from .conftest import (
+    MockEntityFixture,
+    assert_entity_counts,
+    enable_entity,
+    time_changed,
+)
 
 
 @pytest.fixture(name="camera")
@@ -60,14 +65,10 @@ async def camera_fixture(
         camera_obj.id: camera_obj,
     }
 
-    with patch("homeassistant.components.unifiprotect.PLATFORMS", [Platform.CAMERA]):
+    await hass.config_entries.async_setup(mock_entry.entry.entry_id)
-        await hass.config_entries.async_setup(mock_entry.entry.entry_id)
+    await hass.async_block_till_done()
-        await hass.async_block_till_done()
 
-    entity_registry = er.async_get(hass)
+    assert_entity_counts(hass, Platform.CAMERA, 2, 1)
-
-    assert len(hass.states.async_all()) == 1
-    assert len(entity_registry.entities) == 2
 
     yield (camera_obj, "camera.test_camera_high")
 
@@ -264,15 +265,10 @@ async def test_basic_setup(
         camera_all_channels.id: camera_all_channels,
         camera_no_channels.id: camera_no_channels,
     }
+    await hass.config_entries.async_setup(mock_entry.entry.entry_id)
+    await hass.async_block_till_done()
 
-    with patch("homeassistant.components.unifiprotect.PLATFORMS", [Platform.CAMERA]):
+    assert_entity_counts(hass, Platform.CAMERA, 11, 4)
-        await hass.config_entries.async_setup(mock_entry.entry.entry_id)
-        await hass.async_block_till_done()
-
-    entity_registry = er.async_get(hass)
-
-    assert len(hass.states.async_all()) == 4
-    assert len(entity_registry.entities) == 11
 
     # test camera 1
     entity_id = validate_default_camera_entity(hass, camera_high_only, 0)

tests/components/unifiprotect/test_light.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from copy import copy
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 from pyunifiprotect.data import Light
@@ -20,7 +20,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from .conftest import MockEntityFixture
+from .conftest import MockEntityFixture, assert_entity_counts
 
 
 @pytest.fixture(name="light")
@@ -41,14 +41,10 @@ async def light_fixture(
         light_obj.id: light_obj,
     }
 
-    with patch("homeassistant.components.unifiprotect.PLATFORMS", [Platform.LIGHT]):
+    await hass.config_entries.async_setup(mock_entry.entry.entry_id)
-        await hass.config_entries.async_setup(mock_entry.entry.entry_id)
+    await hass.async_block_till_done()
-        await hass.async_block_till_done()
 
-    entity_registry = er.async_get(hass)
+    assert_entity_counts(hass, Platform.LIGHT, 1, 1)
-
-    assert len(hass.states.async_all()) == 1
-    assert len(entity_registry.entities) == 1
 
     yield (light_obj, "light.test_light")
 

tests/components/unifiprotect/test_media_player.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from copy import copy
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 from pyunifiprotect.data import Camera
@@ -26,7 +26,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
 
-from .conftest import MockEntityFixture
+from .conftest import MockEntityFixture, assert_entity_counts
 
 
 @pytest.fixture(name="camera")
@@ -50,16 +50,10 @@ async def camera_fixture(
         camera_obj.id: camera_obj,
     }
 
-    with patch(
+    await hass.config_entries.async_setup(mock_entry.entry.entry_id)
-        "homeassistant.components.unifiprotect.PLATFORMS", [Platform.MEDIA_PLAYER]
+    await hass.async_block_till_done()
-    ):
-        await hass.config_entries.async_setup(mock_entry.entry.entry_id)
-        await hass.async_block_till_done()
 
-    entity_registry = er.async_get(hass)
+    assert_entity_counts(hass, Platform.MEDIA_PLAYER, 1, 1)
-
-    assert len(hass.states.async_all()) == 1
-    assert len(entity_registry.entities) == 1
 
     yield (camera_obj, "media_player.test_camera_speaker")