diff --git a/tests/components/unifiprotect/conftest.py b/tests/components/unifiprotect/conftest.py index b2f3315e095..1b547b179e2 100644 --- a/tests/components/unifiprotect/conftest.py +++ b/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 diff --git a/tests/components/unifiprotect/test_button.py b/tests/components/unifiprotect/test_button.py index 931765dd0bf..ec267a8c68f 100644 --- a/tests/components/unifiprotect/test_button.py +++ b/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.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 len(hass.states.async_all()) == 0 - assert len(entity_registry.entities) == 1 + assert_entity_counts(hass, Platform.BUTTON, 1, 0) yield (camera_obj, "button.test_camera_reboot_device") diff --git a/tests/components/unifiprotect/test_camera.py b/tests/components/unifiprotect/test_camera.py index eb1c640f8a1..ec0b4c96948 100644 --- a/tests/components/unifiprotect/test_camera.py +++ b/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.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 len(hass.states.async_all()) == 1 - assert len(entity_registry.entities) == 2 + assert_entity_counts(hass, Platform.CAMERA, 2, 1) 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]): - 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 + assert_entity_counts(hass, Platform.CAMERA, 11, 4) # test camera 1 entity_id = validate_default_camera_entity(hass, camera_high_only, 0) diff --git a/tests/components/unifiprotect/test_light.py b/tests/components/unifiprotect/test_light.py index c26b00f6da4..8154b5435b6 100644 --- a/tests/components/unifiprotect/test_light.py +++ b/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.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 len(hass.states.async_all()) == 1 - assert len(entity_registry.entities) == 1 + assert_entity_counts(hass, Platform.LIGHT, 1, 1) yield (light_obj, "light.test_light") diff --git a/tests/components/unifiprotect/test_media_player.py b/tests/components/unifiprotect/test_media_player.py index fc39b5ec805..844f01b34c1 100644 --- a/tests/components/unifiprotect/test_media_player.py +++ b/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( - "homeassistant.components.unifiprotect.PLATFORMS", [Platform.MEDIA_PLAYER] - ): - await hass.config_entries.async_setup(mock_entry.entry.entry_id) - 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 len(hass.states.async_all()) == 1 - assert len(entity_registry.entities) == 1 + assert_entity_counts(hass, Platform.MEDIA_PLAYER, 1, 1) yield (camera_obj, "media_player.test_camera_speaker")