diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py
index 4e460bab8f8..798b613b18d 100644
--- a/tests/components/unifi/conftest.py
+++ b/tests/components/unifi/conftest.py
@@ -3,10 +3,10 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Coroutine, Generator
 from datetime import timedelta
 from types import MappingProxyType
-from typing import Any
+from typing import Any, Protocol
 from unittest.mock import AsyncMock, patch
 
 from aiounifi.models.message import MessageKey
@@ -16,7 +16,6 @@ import pytest
 from homeassistant.components.unifi import STORAGE_KEY, STORAGE_VERSION
 from homeassistant.components.unifi.const import CONF_SITE_ID, DOMAIN as UNIFI_DOMAIN
 from homeassistant.components.unifi.hub.websocket import RETRY_TIMER
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -52,6 +51,20 @@ CONTROLLER_HOST = {
     "uptime": 1562600160,
 }
 
+type ConfigEntryFactoryType = Callable[[], Coroutine[Any, Any, MockConfigEntry]]
+
+
+class WebsocketMessageMock(Protocol):
+    """Fixture to mock websocket message."""
+
+    def __call__(
+        self,
+        *,
+        message: MessageKey | None = None,
+        data: list[dict[str, Any]] | dict[str, Any] | None = None,
+    ) -> None:
+        """Send websocket message."""
+
 
 @pytest.fixture(autouse=True, name="mock_discovery")
 def fixture_discovery():
@@ -96,7 +109,7 @@ def fixture_config_entry(
     hass: HomeAssistant,
     config_entry_data: MappingProxyType[str, Any],
     config_entry_options: MappingProxyType[str, Any],
-) -> ConfigEntry:
+) -> MockConfigEntry:
     """Define a config entry fixture."""
     config_entry = MockConfigEntry(
         domain=UNIFI_DOMAIN,
@@ -295,12 +308,12 @@ def fixture_default_requests(
 @pytest.fixture(name="config_entry_factory")
 async def fixture_config_entry_factory(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: MockConfigEntry,
     mock_requests: Callable[[str, str], None],
-) -> Callable[[], ConfigEntry]:
+) -> ConfigEntryFactoryType:
     """Fixture factory that can set up UniFi network integration."""
 
-    async def __mock_setup_config_entry() -> ConfigEntry:
+    async def __mock_setup_config_entry() -> MockConfigEntry:
         mock_requests(config_entry.data[CONF_HOST], config_entry.data[CONF_SITE_ID])
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
@@ -311,8 +324,8 @@ async def fixture_config_entry_factory(
 
 @pytest.fixture(name="config_entry_setup")
 async def fixture_config_entry_setup(
-    config_entry_factory: Callable[[], ConfigEntry],
-) -> ConfigEntry:
+    config_entry_factory: ConfigEntryFactoryType,
+) -> MockConfigEntry:
     """Fixture providing a set up instance of UniFi network integration."""
     return await config_entry_factory()
 
@@ -382,13 +395,15 @@ def fixture_aiounifi_websocket_state(
 
 
 @pytest.fixture(name="mock_websocket_message")
-def fixture_aiounifi_websocket_message(_mock_websocket: AsyncMock):
+def fixture_aiounifi_websocket_message(
+    _mock_websocket: AsyncMock,
+) -> WebsocketMessageMock:
     """No real websocket allowed."""
 
     def make_websocket_call(
         *,
         message: MessageKey | None = None,
-        data: list[dict] | dict | None = None,
+        data: list[dict[str, Any]] | dict[str, Any] | None = None,
     ) -> None:
         """Generate a websocket call."""
         message_handler = _mock_websocket.call_args[0][0]
diff --git a/tests/components/unifi/test_button.py b/tests/components/unifi/test_button.py
index 9af96b64a50..fc3aeccea9f 100644
--- a/tests/components/unifi/test_button.py
+++ b/tests/components/unifi/test_button.py
@@ -11,7 +11,7 @@ from syrupy import SnapshotAssertion
 
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN
 from homeassistant.components.unifi.const import CONF_SITE_ID
-from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY, ConfigEntry
+from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
 from homeassistant.const import (
     CONF_HOST,
     CONTENT_TYPE_JSON,
@@ -23,7 +23,13 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 import homeassistant.util.dt as dt_util
 
-from tests.common import async_fire_time_changed, snapshot_platform
+from .conftest import (
+    ConfigEntryFactoryType,
+    WebsocketMessageMock,
+    WebsocketStateManager,
+)
+
+from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 RANDOM_TOKEN = "random_token"
@@ -134,7 +140,7 @@ WLAN_REGENERATE_PASSWORD = [
 async def test_entity_and_device_data(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    config_entry_factory,
+    config_entry_factory: ConfigEntryFactoryType,
     site_payload: dict[str, Any],
     snapshot: SnapshotAssertion,
 ) -> None:
@@ -150,8 +156,8 @@ async def test_entity_and_device_data(
 async def _test_button_entity(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_state,
-    config_entry: ConfigEntry,
+    mock_websocket_state: WebsocketStateManager,
+    config_entry: MockConfigEntry,
     entity_id: str,
     request_method: str,
     request_path: str,
@@ -221,8 +227,8 @@ async def _test_button_entity(
 async def test_device_button_entities(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
-    mock_websocket_state,
+    config_entry_setup: MockConfigEntry,
+    mock_websocket_state: WebsocketStateManager,
     entity_id: str,
     request_method: str,
     request_path: str,
@@ -269,8 +275,8 @@ async def test_wlan_button_entities(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
-    mock_websocket_state,
+    config_entry_setup: MockConfigEntry,
+    mock_websocket_state: WebsocketStateManager,
     entity_id: str,
     request_method: str,
     request_path: str,
@@ -308,7 +314,7 @@ async def test_wlan_button_entities(
 @pytest.mark.usefixtures("config_entry_setup")
 async def test_power_cycle_availability(
     hass: HomeAssistant,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: dict[str, Any],
 ) -> None:
     """Verify that disabling PoE marks entity as unavailable."""
diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py
index fc0d2626eb6..1d745511dc5 100644
--- a/tests/components/unifi/test_config_flow.py
+++ b/tests/components/unifi/test_config_flow.py
@@ -1,6 +1,5 @@
 """Test UniFi Network config flow."""
 
-from collections.abc import Callable
 import socket
 from unittest.mock import PropertyMock, patch
 
@@ -25,7 +24,7 @@ from homeassistant.components.unifi.const import (
     CONF_TRACK_WIRED_CLIENTS,
     DOMAIN as UNIFI_DOMAIN,
 )
-from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntry
+from homeassistant.config_entries import SOURCE_REAUTH
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -36,6 +35,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
+from .conftest import ConfigEntryFactoryType
+
 from tests.common import MockConfigEntry
 
 CLIENTS = [{"mac": "00:00:00:00:00:01"}]
@@ -296,7 +297,7 @@ async def test_flow_fails_hub_unavailable(hass: HomeAssistant) -> None:
 
 
 async def test_reauth_flow_update_configuration(
-    hass: HomeAssistant, config_entry_setup: ConfigEntry
+    hass: HomeAssistant, config_entry_setup: MockConfigEntry
 ) -> None:
     """Verify reauth flow can update hub configuration."""
     config_entry = config_entry_setup
@@ -337,7 +338,7 @@ async def test_reauth_flow_update_configuration(
 
 
 async def test_reauth_flow_update_configuration_on_not_loaded_entry(
-    hass: HomeAssistant, config_entry_factory: Callable[[], ConfigEntry]
+    hass: HomeAssistant, config_entry_factory: ConfigEntryFactoryType
 ) -> None:
     """Verify reauth flow can update hub configuration on a not loaded entry."""
     with patch("aiounifi.Controller.login", side_effect=aiounifi.errors.RequestError):
@@ -379,7 +380,7 @@ async def test_reauth_flow_update_configuration_on_not_loaded_entry(
 @pytest.mark.parametrize("wlan_payload", [WLANS])
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
 async def test_advanced_option_flow(
-    hass: HomeAssistant, config_entry_setup: ConfigEntry
+    hass: HomeAssistant, config_entry_setup: MockConfigEntry
 ) -> None:
     """Test advanced config flow options."""
     config_entry = config_entry_setup
@@ -463,7 +464,7 @@ async def test_advanced_option_flow(
 
 @pytest.mark.parametrize("client_payload", [CLIENTS])
 async def test_simple_option_flow(
-    hass: HomeAssistant, config_entry_setup: ConfigEntry
+    hass: HomeAssistant, config_entry_setup: MockConfigEntry
 ) -> None:
     """Test simple config flow options."""
     config_entry = config_entry_setup
diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py
index f2480a4f050..c653370656d 100644
--- a/tests/components/unifi/test_device_tracker.py
+++ b/tests/components/unifi/test_device_tracker.py
@@ -1,6 +1,5 @@
 """The tests for the UniFi Network device tracker platform."""
 
-from collections.abc import Callable
 from datetime import timedelta
 from types import MappingProxyType
 from typing import Any
@@ -24,13 +23,18 @@ from homeassistant.components.unifi.const import (
     DEFAULT_DETECTION_TIME,
     DOMAIN as UNIFI_DOMAIN,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_HOME, STATE_NOT_HOME, STATE_UNAVAILABLE, Platform
 from homeassistant.core import HomeAssistant, State
 from homeassistant.helpers import entity_registry as er
 import homeassistant.util.dt as dt_util
 
-from tests.common import async_fire_time_changed, snapshot_platform
+from .conftest import (
+    ConfigEntryFactoryType,
+    WebsocketMessageMock,
+    WebsocketStateManager,
+)
+
+from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 
 WIRED_CLIENT_1 = {
     "hostname": "wd_client_1",
@@ -96,7 +100,7 @@ SWITCH_1 = {
 async def test_entity_and_device_data(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    config_entry_factory,
+    config_entry_factory: ConfigEntryFactoryType,
     snapshot: SnapshotAssertion,
 ) -> None:
     """Validate entity and device data with and without admin rights."""
@@ -112,8 +116,8 @@ async def test_entity_and_device_data(
 @pytest.mark.usefixtures("mock_device_registry")
 async def test_client_state_update(
     hass: HomeAssistant,
-    mock_websocket_message,
-    config_entry_factory: Callable[[], ConfigEntry],
+    mock_websocket_message: WebsocketMessageMock,
+    config_entry_factory: ConfigEntryFactoryType,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify tracking of wireless clients."""
@@ -165,7 +169,7 @@ async def test_client_state_update(
 async def test_client_state_from_event_source(
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify update state of client based on event source."""
@@ -247,8 +251,8 @@ async def test_client_state_from_event_source(
 async def test_tracked_device_state_change(
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,
-    config_entry_factory: Callable[[], ConfigEntry],
-    mock_websocket_message,
+    config_entry_factory: ConfigEntryFactoryType,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
     state: int,
     interval: int,
@@ -289,7 +293,9 @@ async def test_tracked_device_state_change(
 @pytest.mark.usefixtures("config_entry_setup")
 @pytest.mark.usefixtures("mock_device_registry")
 async def test_remove_clients(
-    hass: HomeAssistant, mock_websocket_message, client_payload: list[dict[str, Any]]
+    hass: HomeAssistant,
+    mock_websocket_message: WebsocketMessageMock,
+    client_payload: list[dict[str, Any]],
 ) -> None:
     """Test the remove_items function with some clients."""
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
@@ -309,7 +315,10 @@ async def test_remove_clients(
 @pytest.mark.parametrize("device_payload", [[SWITCH_1]])
 @pytest.mark.usefixtures("config_entry_setup")
 @pytest.mark.usefixtures("mock_device_registry")
-async def test_hub_state_change(hass: HomeAssistant, mock_websocket_state) -> None:
+async def test_hub_state_change(
+    hass: HomeAssistant,
+    mock_websocket_state: WebsocketStateManager,
+) -> None:
     """Verify entities state reflect on hub connection becoming unavailable."""
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
     assert hass.states.get("device_tracker.ws_client_1").state == STATE_NOT_HOME
@@ -330,7 +339,7 @@ async def test_hub_state_change(hass: HomeAssistant, mock_websocket_state) -> No
 async def test_option_ssid_filter(
     hass: HomeAssistant,
     mock_websocket_message,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Test the SSID filter works.
@@ -434,7 +443,7 @@ async def test_option_ssid_filter(
 async def test_wireless_client_go_wired_issue(
     hass: HomeAssistant,
     mock_websocket_message,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Test the solution to catch wireless device go wired UniFi issue.
@@ -494,7 +503,7 @@ async def test_wireless_client_go_wired_issue(
 async def test_option_ignore_wired_bug(
     hass: HomeAssistant,
     mock_websocket_message,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Test option to ignore wired bug."""
@@ -571,8 +580,8 @@ async def test_option_ignore_wired_bug(
 async def test_restoring_client(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    config_entry: ConfigEntry,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry: MockConfigEntry,
+    config_entry_factory: ConfigEntryFactoryType,
     client_payload: list[dict[str, Any]],
     clients_all_payload: list[dict[str, Any]],
 ) -> None:
@@ -645,7 +654,7 @@ async def test_restoring_client(
 @pytest.mark.usefixtures("mock_device_registry")
 async def test_config_entry_options_track(
     hass: HomeAssistant,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     config_entry_options: MappingProxyType[str, Any],
     counts: tuple[int],
     expected: tuple[tuple[bool | None, ...], ...],
diff --git a/tests/components/unifi/test_diagnostics.py b/tests/components/unifi/test_diagnostics.py
index 3963de2deb3..80359a9c75c 100644
--- a/tests/components/unifi/test_diagnostics.py
+++ b/tests/components/unifi/test_diagnostics.py
@@ -9,9 +9,9 @@ from homeassistant.components.unifi.const import (
     CONF_ALLOW_UPTIME_SENSORS,
     CONF_BLOCK_CLIENT,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
+from tests.common import MockConfigEntry
 from tests.components.diagnostics import get_diagnostics_for_config_entry
 from tests.typing import ClientSessionGenerator
 
@@ -122,7 +122,7 @@ DPI_GROUP_DATA = [
 async def test_entry_diagnostics(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     snapshot: SnapshotAssertion,
 ) -> None:
     """Test config entry diagnostics."""
diff --git a/tests/components/unifi/test_hub.py b/tests/components/unifi/test_hub.py
index 0d75a83c5f5..af134c7449b 100644
--- a/tests/components/unifi/test_hub.py
+++ b/tests/components/unifi/test_hub.py
@@ -1,6 +1,5 @@
 """Test UniFi Network."""
 
-from collections.abc import Callable
 from http import HTTPStatus
 from types import MappingProxyType
 from typing import Any
@@ -12,18 +11,21 @@ import pytest
 from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
 from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
 from homeassistant.components.unifi.hub import get_unifi_api
-from homeassistant.config_entries import ConfigEntry, ConfigEntryState
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 import homeassistant.util.dt as dt_util
 
+from .conftest import ConfigEntryFactoryType, WebsocketStateManager
+
+from tests.common import MockConfigEntry
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 
 async def test_hub_setup(
     device_registry: dr.DeviceRegistry,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
 ) -> None:
     """Successful setup."""
     with patch(
@@ -54,7 +56,7 @@ async def test_hub_setup(
 
 
 async def test_reset_after_successful_setup(
-    hass: HomeAssistant, config_entry_setup: ConfigEntry
+    hass: HomeAssistant, config_entry_setup: MockConfigEntry
 ) -> None:
     """Calling reset when the entry has been setup."""
     assert config_entry_setup.state is ConfigEntryState.LOADED
@@ -64,7 +66,7 @@ async def test_reset_after_successful_setup(
 
 
 async def test_reset_fails(
-    hass: HomeAssistant, config_entry_setup: ConfigEntry
+    hass: HomeAssistant, config_entry_setup: MockConfigEntry
 ) -> None:
     """Calling reset when the entry has been setup can return false."""
     assert config_entry_setup.state is ConfigEntryState.LOADED
@@ -80,8 +82,8 @@ async def test_reset_fails(
 @pytest.mark.usefixtures("mock_device_registry")
 async def test_connection_state_signalling(
     hass: HomeAssistant,
-    mock_websocket_state,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
+    mock_websocket_state: WebsocketStateManager,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify connection statesignalling and connection state are working."""
@@ -110,8 +112,8 @@ async def test_connection_state_signalling(
 
 async def test_reconnect_mechanism(
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_state,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
+    mock_websocket_state: WebsocketStateManager,
 ) -> None:
     """Verify reconnect prints only on first reconnection try."""
     aioclient_mock.clear_requests()
@@ -140,7 +142,10 @@ async def test_reconnect_mechanism(
     ],
 )
 @pytest.mark.usefixtures("config_entry_setup")
-async def test_reconnect_mechanism_exceptions(mock_websocket_state, exception) -> None:
+async def test_reconnect_mechanism_exceptions(
+    mock_websocket_state: WebsocketStateManager,
+    exception: Exception,
+) -> None:
     """Verify async_reconnect calls expected methods."""
     with (
         patch("aiounifi.Controller.login", side_effect=exception),
@@ -170,8 +175,8 @@ async def test_reconnect_mechanism_exceptions(mock_websocket_state, exception) -
 )
 async def test_get_unifi_api_fails_to_connect(
     hass: HomeAssistant,
-    side_effect,
-    raised_exception,
+    side_effect: Exception,
+    raised_exception: Exception,
     config_entry_data: MappingProxyType[str, Any],
 ) -> None:
     """Check that get_unifi_api can handle UniFi Network being unavailable."""
diff --git a/tests/components/unifi/test_image.py b/tests/components/unifi/test_image.py
index 6733845c52f..dc37d7cb8b7 100644
--- a/tests/components/unifi/test_image.py
+++ b/tests/components/unifi/test_image.py
@@ -18,6 +18,12 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 from homeassistant.util import dt as dt_util
 
+from .conftest import (
+    ConfigEntryFactoryType,
+    WebsocketMessageMock,
+    WebsocketStateManager,
+)
+
 from tests.common import async_fire_time_changed, snapshot_platform
 from tests.typing import ClientSessionGenerator
 
@@ -82,7 +88,7 @@ WLAN = {
 async def test_entity_and_device_data(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    config_entry_factory,
+    config_entry_factory: ConfigEntryFactoryType,
     site_payload: dict[str, Any],
     snapshot: SnapshotAssertion,
 ) -> None:
@@ -102,7 +108,7 @@ async def test_wlan_qr_code(
     entity_registry: er.EntityRegistry,
     hass_client: ClientSessionGenerator,
     snapshot: SnapshotAssertion,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
 ) -> None:
     """Test the update_clients function when no clients are found."""
     assert len(hass.states.async_entity_ids(IMAGE_DOMAIN)) == 0
@@ -151,7 +157,9 @@ async def test_wlan_qr_code(
 @pytest.mark.parametrize("wlan_payload", [[WLAN]])
 @pytest.mark.usefixtures("config_entry_setup")
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
-async def test_hub_state_change(hass: HomeAssistant, mock_websocket_state) -> None:
+async def test_hub_state_change(
+    hass: HomeAssistant, mock_websocket_state: WebsocketStateManager
+) -> None:
     """Verify entities state reflect on hub becoming unavailable."""
     assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE
 
@@ -167,7 +175,9 @@ async def test_hub_state_change(hass: HomeAssistant, mock_websocket_state) -> No
 @pytest.mark.parametrize("wlan_payload", [[WLAN]])
 @pytest.mark.usefixtures("config_entry_setup")
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
-async def test_source_availability(hass: HomeAssistant, mock_websocket_message) -> None:
+async def test_source_availability(
+    hass: HomeAssistant, mock_websocket_message: WebsocketMessageMock
+) -> None:
     """Verify entities state reflect on source becoming unavailable."""
     assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE
 
diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py
index de08ba2c6d7..68f80555cd6 100644
--- a/tests/components/unifi/test_init.py
+++ b/tests/components/unifi/test_init.py
@@ -1,6 +1,5 @@
 """Test UniFi Network integration setup process."""
 
-from collections.abc import Callable
 from typing import Any
 from unittest.mock import patch
 
@@ -15,19 +14,23 @@ from homeassistant.components.unifi.const import (
     CONF_TRACK_DEVICES,
 )
 from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
-from homeassistant.config_entries import ConfigEntry, ConfigEntryState
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 from homeassistant.setup import async_setup_component
 
-from .conftest import DEFAULT_CONFIG_ENTRY_ID
+from .conftest import (
+    DEFAULT_CONFIG_ENTRY_ID,
+    ConfigEntryFactoryType,
+    WebsocketMessageMock,
+)
 
 from tests.common import flush_store
 from tests.typing import WebSocketGenerator
 
 
 async def test_setup_entry_fails_config_entry_not_ready(
-    config_entry_factory: Callable[[], ConfigEntry],
+    hass: HomeAssistant, config_entry_factory: ConfigEntryFactoryType
 ) -> None:
     """Failed authentication trigger a reauthentication flow."""
     with patch(
@@ -40,7 +43,7 @@ async def test_setup_entry_fails_config_entry_not_ready(
 
 
 async def test_setup_entry_fails_trigger_reauth_flow(
-    hass: HomeAssistant, config_entry_factory: Callable[[], ConfigEntry]
+    hass: HomeAssistant, config_entry_factory: ConfigEntryFactoryType
 ) -> None:
     """Failed authentication trigger a reauthentication flow."""
     with (
@@ -78,7 +81,7 @@ async def test_setup_entry_fails_trigger_reauth_flow(
 async def test_wireless_clients(
     hass: HomeAssistant,
     hass_storage: dict[str, Any],
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
 ) -> None:
     """Verify wireless clients class."""
     hass_storage[unifi.STORAGE_KEY] = {
@@ -163,10 +166,10 @@ async def test_wireless_clients(
 async def test_remove_config_entry_device(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
     client_payload: list[dict[str, Any]],
     device_payload: list[dict[str, Any]],
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Verify removing a device manually."""
diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py
index 779df6660f0..5af4b297847 100644
--- a/tests/components/unifi/test_sensor.py
+++ b/tests/components/unifi/test_sensor.py
@@ -1,6 +1,5 @@
 """UniFi Network sensor platform tests."""
 
-from collections.abc import Callable
 from copy import deepcopy
 from datetime import datetime, timedelta
 from types import MappingProxyType
@@ -29,7 +28,7 @@ from homeassistant.components.unifi.const import (
     DEFAULT_DETECTION_TIME,
     DEVICE_STATES,
 )
-from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY, ConfigEntry
+from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
     ATTR_FRIENDLY_NAME,
@@ -42,7 +41,13 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 import homeassistant.util.dt as dt_util
 
-from tests.common import async_fire_time_changed
+from .conftest import (
+    ConfigEntryFactoryType,
+    WebsocketMessageMock,
+    WebsocketStateManager,
+)
+
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 DEVICE_1 = {
     "board_rev": 2,
@@ -362,9 +367,9 @@ async def test_no_clients(hass: HomeAssistant) -> None:
 )
 async def test_bandwidth_sensors(
     hass: HomeAssistant,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     config_entry_options: MappingProxyType[str, Any],
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify that bandwidth sensors are working as expected."""
@@ -491,7 +496,9 @@ async def test_bandwidth_sensors(
 @pytest.mark.usefixtures("config_entry_setup")
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
 async def test_remove_sensors(
-    hass: HomeAssistant, mock_websocket_message, client_payload: list[dict[str, Any]]
+    hass: HomeAssistant,
+    mock_websocket_message: WebsocketMessageMock,
+    client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify removing of clients work as expected."""
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 6
@@ -520,8 +527,8 @@ async def test_remove_sensors(
 async def test_poe_port_switches(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
-    mock_websocket_state,
+    mock_websocket_message: WebsocketMessageMock,
+    mock_websocket_state: WebsocketStateManager,
 ) -> None:
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2
@@ -594,9 +601,9 @@ async def test_poe_port_switches(
 async def test_wlan_client_sensors(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
-    mock_websocket_state,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
+    mock_websocket_message: WebsocketMessageMock,
+    mock_websocket_state: WebsocketStateManager,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify that WLAN client sensors are working as expected."""
@@ -736,7 +743,7 @@ async def test_wlan_client_sensors(
 async def test_outlet_power_readings(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
     entity_id: str,
     expected_unique_id: str,
@@ -798,7 +805,7 @@ async def test_outlet_power_readings(
 async def test_device_temperature(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
 ) -> None:
     """Verify that temperature sensors are working as expected."""
@@ -847,7 +854,7 @@ async def test_device_temperature(
 async def test_device_state(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
 ) -> None:
     """Verify that state sensors are working as expected."""
@@ -884,7 +891,7 @@ async def test_device_state(
 async def test_device_system_stats(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
 ) -> None:
     """Verify that device stats sensors are working as expected."""
@@ -979,9 +986,9 @@ async def test_device_system_stats(
 async def test_bandwidth_port_sensors(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     config_entry_options: MappingProxyType[str, Any],
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
 ) -> None:
     """Verify that port bandwidth sensors are working as expected."""
@@ -1096,9 +1103,9 @@ async def test_bandwidth_port_sensors(
 async def test_device_client_sensors(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    config_entry_factory,
-    mock_websocket_message,
-    client_payload,
+    config_entry_factory: ConfigEntryFactoryType,
+    mock_websocket_message: WebsocketMessageMock,
+    client_payload: dict[str, Any],
 ) -> None:
     """Verify that WLAN client sensors are working as expected."""
     client_payload += [
@@ -1246,8 +1253,8 @@ async def test_sensor_sources(
 async def _test_uptime_entity(
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,
-    mock_websocket_message,
-    config_entry_factory: Callable[[], ConfigEntry],
+    mock_websocket_message: WebsocketMessageMock,
+    config_entry_factory: ConfigEntryFactoryType,
     payload: dict[str, Any],
     entity_id: str,
     message_key: MessageKey,
@@ -1326,9 +1333,9 @@ async def test_client_uptime(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     freezer: FrozenDateTimeFactory,
-    mock_websocket_message,
     config_entry_options: MappingProxyType[str, Any],
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
+    mock_websocket_message: WebsocketMessageMock,
     client_payload: list[dict[str, Any]],
     initial_uptime,
     event_uptime,
@@ -1401,8 +1408,8 @@ async def test_device_uptime(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     freezer: FrozenDateTimeFactory,
-    mock_websocket_message,
-    config_entry_factory: Callable[[], ConfigEntry],
+    config_entry_factory: ConfigEntryFactoryType,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
 ) -> None:
     """Verify that device uptime sensors are working as expected."""
@@ -1502,7 +1509,7 @@ async def test_device_uptime(
 async def test_wan_monitor_latency(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    mock_websocket_message,
+    mock_websocket_message: WebsocketMessageMock,
     device_payload: list[dict[str, Any]],
     entity_id: str,
     state: str,
diff --git a/tests/components/unifi/test_services.py b/tests/components/unifi/test_services.py
index bf7058e28ff..a7968a92e22 100644
--- a/tests/components/unifi/test_services.py
+++ b/tests/components/unifi/test_services.py
@@ -10,11 +10,11 @@ from homeassistant.components.unifi.services import (
     SERVICE_RECONNECT_CLIENT,
     SERVICE_REMOVE_CLIENTS,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_DEVICE_ID, CONF_HOST
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 
+from tests.common import MockConfigEntry
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 
@@ -25,7 +25,7 @@ async def test_reconnect_client(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify call to reconnect client is performed as expected."""
@@ -69,7 +69,7 @@ async def test_reconnect_device_without_mac(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
 ) -> None:
     """Verify no call is made if device does not have a known mac."""
     aioclient_mock.clear_requests()
@@ -95,7 +95,7 @@ async def test_reconnect_client_hub_unavailable(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify no call is made if hub is unavailable."""
@@ -127,7 +127,7 @@ async def test_reconnect_client_unknown_mac(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
 ) -> None:
     """Verify no call is made if trying to reconnect a mac unknown to hub."""
     aioclient_mock.clear_requests()
@@ -152,7 +152,7 @@ async def test_reconnect_wired_client(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     client_payload: list[dict[str, Any]],
 ) -> None:
     """Verify no call is made if client is wired."""
@@ -204,7 +204,7 @@ async def test_reconnect_wired_client(
 async def test_remove_clients(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
 ) -> None:
     """Verify removing different variations of clients work."""
     aioclient_mock.clear_requests()
@@ -288,7 +288,7 @@ async def test_services_handle_unloaded_config_entry(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
     device_registry: dr.DeviceRegistry,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     clients_all_payload: dict[str, Any],
 ) -> None:
     """Verify no call is made if config entry is unloaded."""
diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py
index daf64301c8e..6d85437a244 100644
--- a/tests/components/unifi/test_switch.py
+++ b/tests/components/unifi/test_switch.py
@@ -1,6 +1,5 @@
 """UniFi Network switch platform tests."""
 
-from collections.abc import Callable
 from copy import deepcopy
 from datetime import timedelta
 from typing import Any
@@ -22,7 +21,7 @@ from homeassistant.components.unifi.const import (
     CONF_TRACK_DEVICES,
     DOMAIN as UNIFI_DOMAIN,
 )
-from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY, ConfigEntry
+from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
     ATTR_ENTITY_ID,
@@ -37,9 +36,14 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 from homeassistant.util import dt as dt_util
 
-from .conftest import CONTROLLER_HOST
+from .conftest import (
+    CONTROLLER_HOST,
+    ConfigEntryFactoryType,
+    WebsocketMessageMock,
+    WebsocketStateManager,
+)
 
-from tests.common import async_fire_time_changed
+from tests.common import MockConfigEntry, async_fire_time_changed
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 CLIENT_1 = {
@@ -841,12 +845,11 @@ async def test_not_admin(hass: HomeAssistant) -> None:
 @pytest.mark.parametrize("clients_all_payload", [[BLOCKED, UNBLOCKED, CLIENT_1]])
 @pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
-@pytest.mark.usefixtures("config_entry_setup")
 async def test_switches(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
 ) -> None:
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3
@@ -930,7 +933,9 @@ async def test_switches(
 @pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
 @pytest.mark.usefixtures("config_entry_setup")
-async def test_remove_switches(hass: HomeAssistant, mock_websocket_message) -> None:
+async def test_remove_switches(
+    hass: HomeAssistant, mock_websocket_message: WebsocketMessageMock
+) -> None:
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
 
@@ -967,8 +972,8 @@ async def test_remove_switches(hass: HomeAssistant, mock_websocket_message) -> N
 async def test_block_switches(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message,
-    config_entry_setup: ConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
+    config_entry_setup: MockConfigEntry,
 ) -> None:
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
@@ -1027,7 +1032,9 @@ async def test_block_switches(
 @pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
 @pytest.mark.usefixtures("config_entry_setup")
-async def test_dpi_switches(hass: HomeAssistant, mock_websocket_message) -> None:
+async def test_dpi_switches(
+    hass: HomeAssistant, mock_websocket_message: WebsocketMessageMock
+) -> None:
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
@@ -1053,7 +1060,7 @@ async def test_dpi_switches(hass: HomeAssistant, mock_websocket_message) -> None
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
 @pytest.mark.usefixtures("config_entry_setup")
 async def test_dpi_switches_add_second_app(
-    hass: HomeAssistant, mock_websocket_message
+    hass: HomeAssistant, mock_websocket_message: WebsocketMessageMock
 ) -> None:
     """Test the update_items function with some clients."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@@ -1104,12 +1111,10 @@ async def test_dpi_switches_add_second_app(
 
 
 @pytest.mark.parametrize(("traffic_rule_payload"), [([TRAFFIC_RULE])])
-@pytest.mark.usefixtures("config_entry_setup")
 async def test_traffic_rules(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
     traffic_rule_payload: list[dict[str, Any]],
 ) -> None:
     """Test control of UniFi traffic rules."""
@@ -1172,8 +1177,8 @@ async def test_traffic_rules(
 async def test_outlet_switches(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message,
-    config_entry_setup: ConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
+    config_entry_setup: MockConfigEntry,
     device_payload: list[dict[str, Any]],
     entity_id: str,
     outlet_index: int,
@@ -1268,7 +1273,7 @@ async def test_outlet_switches(
 )
 @pytest.mark.usefixtures("config_entry_setup")
 async def test_new_client_discovered_on_block_control(
-    hass: HomeAssistant, mock_websocket_message
+    hass: HomeAssistant, mock_websocket_message: WebsocketMessageMock
 ) -> None:
     """Test if 2nd update has a new client."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
@@ -1286,7 +1291,9 @@ async def test_new_client_discovered_on_block_control(
 )
 @pytest.mark.parametrize("clients_all_payload", [[BLOCKED, UNBLOCKED]])
 async def test_option_block_clients(
-    hass: HomeAssistant, config_entry_setup: ConfigEntry, clients_all_payload
+    hass: HomeAssistant,
+    config_entry_setup: MockConfigEntry,
+    clients_all_payload: list[dict[str, Any]],
 ) -> None:
     """Test the changes to option reflects accordingly."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@@ -1334,7 +1341,7 @@ async def test_option_block_clients(
 @pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
 @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
 async def test_option_remove_switches(
-    hass: HomeAssistant, config_entry_setup: ConfigEntry
+    hass: HomeAssistant, config_entry_setup: MockConfigEntry
 ) -> None:
     """Test removal of DPI switch when options updated."""
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@@ -1352,8 +1359,8 @@ async def test_poe_port_switches(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message,
-    config_entry_setup: ConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
+    config_entry_setup: MockConfigEntry,
     device_payload: list[dict[str, Any]],
 ) -> None:
     """Test PoE port entities work."""
@@ -1451,8 +1458,8 @@ async def test_wlan_switches(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message,
-    config_entry_setup: ConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
+    config_entry_setup: MockConfigEntry,
     wlan_payload: list[dict[str, Any]],
 ) -> None:
     """Test control of UniFi WLAN availability."""
@@ -1507,8 +1514,8 @@ async def test_port_forwarding_switches(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
     aioclient_mock: AiohttpClientMocker,
-    mock_websocket_message,
-    config_entry_setup: ConfigEntry,
+    mock_websocket_message: WebsocketMessageMock,
+    config_entry_setup: MockConfigEntry,
     port_forward_payload: list[dict[str, Any]],
 ) -> None:
     """Test control of UniFi port forwarding."""
@@ -1606,9 +1613,9 @@ async def test_port_forwarding_switches(
 async def test_updating_unique_id(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    config_entry_factory: Callable[[], ConfigEntry],
-    config_entry: ConfigEntry,
-    device_payload,
+    config_entry_factory: ConfigEntryFactoryType,
+    config_entry: MockConfigEntry,
+    device_payload: list[dict[str, Any]],
 ) -> None:
     """Verify outlet control and poe control unique ID update works."""
     entity_registry.async_get_or_create(
@@ -1644,7 +1651,9 @@ async def test_updating_unique_id(
 @pytest.mark.parametrize("wlan_payload", [[WLAN]])
 @pytest.mark.usefixtures("config_entry_setup")
 @pytest.mark.usefixtures("entity_registry_enabled_by_default")
-async def test_hub_state_change(hass: HomeAssistant, mock_websocket_state) -> None:
+async def test_hub_state_change(
+    hass: HomeAssistant, mock_websocket_state: WebsocketStateManager
+) -> None:
     """Verify entities state reflect on hub connection becoming unavailable."""
     entity_ids = (
         "switch.block_client_2",
diff --git a/tests/components/unifi/test_update.py b/tests/components/unifi/test_update.py
index a8fe9231159..7bf4b9aec9d 100644
--- a/tests/components/unifi/test_update.py
+++ b/tests/components/unifi/test_update.py
@@ -16,7 +16,6 @@ from homeassistant.components.update import (
     DOMAIN as UPDATE_DOMAIN,
     SERVICE_INSTALL,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_ENTITY_ID,
     CONF_HOST,
@@ -28,7 +27,13 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 
-from tests.common import snapshot_platform
+from .conftest import (
+    ConfigEntryFactoryType,
+    WebsocketMessageMock,
+    WebsocketStateManager,
+)
+
+from tests.common import MockConfigEntry, snapshot_platform
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 # Device with new firmware available
@@ -74,7 +79,7 @@ DEVICE_2 = {
 async def test_entity_and_device_data(
     hass: HomeAssistant,
     entity_registry: er.EntityRegistry,
-    config_entry_factory,
+    config_entry_factory: ConfigEntryFactoryType,
     snapshot: SnapshotAssertion,
 ) -> None:
     """Validate entity and device data with and without admin rights."""
@@ -85,7 +90,9 @@ async def test_entity_and_device_data(
 
 @pytest.mark.parametrize("device_payload", [[DEVICE_1]])
 @pytest.mark.usefixtures("config_entry_setup")
-async def test_device_updates(hass: HomeAssistant, mock_websocket_message) -> None:
+async def test_device_updates(
+    hass: HomeAssistant, mock_websocket_message: WebsocketMessageMock
+) -> None:
     """Test the update_items function with some devices."""
     device_1_state = hass.states.get("update.device_1")
     assert device_1_state.state == STATE_ON
@@ -122,7 +129,7 @@ async def test_device_updates(hass: HomeAssistant, mock_websocket_message) -> No
 async def test_install(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    config_entry_setup: ConfigEntry,
+    config_entry_setup: MockConfigEntry,
 ) -> None:
     """Test the device update install call."""
     device_state = hass.states.get("update.device_1")
@@ -154,7 +161,9 @@ async def test_install(
 
 @pytest.mark.parametrize("device_payload", [[DEVICE_1]])
 @pytest.mark.usefixtures("config_entry_setup")
-async def test_hub_state_change(hass: HomeAssistant, mock_websocket_state) -> None:
+async def test_hub_state_change(
+    hass: HomeAssistant, mock_websocket_state: WebsocketStateManager
+) -> None:
     """Verify entities state reflect on hub becoming unavailable."""
     assert hass.states.get("update.device_1").state == STATE_ON