Improve hass_ws_client type hint in tests (#89703)

tests/components/automation/test_init.py

@@ -1,11 +1,9 @@
 """The tests for the automation component."""
 import asyncio
-from collections.abc import Awaitable, Callable
 from datetime import timedelta
 import logging
 from unittest.mock import Mock, patch
 
-from aiohttp import ClientWebSocketResponse
 import pytest
 
 import homeassistant.components.automation as automation
@@ -1437,7 +1435,7 @@ async def test_automation_bad_config_validation(
 async def test_automation_with_error_in_script(
     hass: HomeAssistant,
     caplog: pytest.LogCaptureFixture,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test automation with an error in script."""
     assert await async_setup_component(

tests/components/backup/test_websocket.py

@@ -1,18 +1,18 @@
 """Tests for the Backup integration."""
-from collections.abc import Awaitable, Callable
 from unittest.mock import patch
 
-from aiohttp import ClientWebSocketResponse
 import pytest
 
 from homeassistant.core import HomeAssistant
 
 from .common import TEST_BACKUP, setup_backup_integration
 
+from tests.typing import WebSocketGenerator
+
 
 async def test_info(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test getting backup info."""
     await setup_backup_integration(hass)
@@ -34,7 +34,7 @@ async def test_info(
 
 async def test_remove(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test removing a backup file."""
@@ -55,7 +55,7 @@ async def test_remove(
 
 async def test_generate(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test removing a backup file."""
     await setup_backup_integration(hass)

tests/components/devolo_home_control/test_init.py

@@ -1,8 +1,6 @@
 """Tests for the devolo Home Control integration."""
-from collections.abc import Awaitable, Callable
 from unittest.mock import patch
 
-from aiohttp import ClientWebSocketResponse
 from devolo_home_control_api.exceptions.gateway import GatewayOfflineError
 import pytest
 
@@ -16,6 +14,8 @@ from homeassistant.setup import async_setup_component
 from . import configure_integration
 from .mocks import HomeControlMock, HomeControlMockBinarySensor
 
+from tests.typing import WebSocketGenerator
+
 
 async def test_setup_entry(hass: HomeAssistant, mock_zeroconf: None) -> None:
     """Test setup entry."""
@@ -64,7 +64,7 @@ async def test_unload_entry(hass: HomeAssistant) -> None:
 
 async def test_remove_device(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ):
     """Test removing a device."""
     assert await async_setup_component(hass, "config", {})

tests/components/dlink/test_switch.py

@@ -1,5 +1,4 @@
 """Switch tests for the D-Link Smart Plug integration."""
-from collections.abc import Awaitable, Callable
 
 from homeassistant.components.dlink import DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
@@ -16,11 +15,12 @@ from homeassistant.setup import async_setup_component
 from .conftest import ComponentSetup
 
 from tests.components.repairs import get_repairs
+from tests.typing import WebSocketGenerator
 
 
 async def test_switch_state(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[None]],
+    hass_ws_client: WebSocketGenerator,
     setup_integration: ComponentSetup,
 ) -> None:
     """Test we get the switch status."""

tests/components/google/test_calendar.py

@@ -8,7 +8,6 @@ from typing import Any
 from unittest.mock import patch
 import urllib
 
-from aiohttp import ClientWebSocketResponse
 from aiohttp.client_exceptions import ClientError
 from gcal_sync.auth import API_BASE_URL
 import pytest
@@ -32,7 +31,7 @@ from .conftest import (
 
 from tests.common import async_fire_time_changed
 from tests.test_util.aiohttp import AiohttpClientMocker
-from tests.typing import ClientSessionGenerator
+from tests.typing import ClientSessionGenerator, WebSocketGenerator
 
 TEST_ENTITY = TEST_API_ENTITY
 TEST_ENTITY_NAME = TEST_API_ENTITY_NAME
@@ -134,7 +133,7 @@ ClientFixture = Callable[[], Awaitable[Client]]
 @pytest.fixture
 async def ws_client(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> ClientFixture:
     """Fixture for creating the test websocket client."""
 

tests/components/matter/test_api.py

@@ -1,8 +1,6 @@
 """Test the api module."""
-from collections.abc import Awaitable, Callable
 from unittest.mock import MagicMock, call
 
-from aiohttp import ClientWebSocketResponse
 from matter_server.common.errors import InvalidCommand, NodeCommissionFailed
 import pytest
 
@@ -10,13 +8,14 @@ from homeassistant.components.matter.api import ID, TYPE
 from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry
+from tests.typing import WebSocketGenerator
 
 
 # This tests needs to be adjusted to remove lingering tasks
 @pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_commission(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     matter_client: MagicMock,
     integration: MockConfigEntry,
 ) -> None:
@@ -58,7 +57,7 @@ async def test_commission(
 @pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_commission_on_network(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     matter_client: MagicMock,
     integration: MockConfigEntry,
 ) -> None:
@@ -100,7 +99,7 @@ async def test_commission_on_network(
 @pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_set_thread_dataset(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     matter_client: MagicMock,
     integration: MockConfigEntry,
 ) -> None:
@@ -142,7 +141,7 @@ async def test_set_thread_dataset(
 @pytest.mark.parametrize("expected_lingering_tasks", [True])
 async def test_set_wifi_credentials(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     matter_client: MagicMock,
     integration: MockConfigEntry,
 ) -> None:

tests/components/matter/test_init.py

@@ -2,10 +2,9 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Awaitable, Callable, Generator
+from collections.abc import Generator
 from unittest.mock import AsyncMock, MagicMock, call, patch
 
-from aiohttp import ClientWebSocketResponse
 from matter_server.client.exceptions import CannotConnect, InvalidServerVersion
 from matter_server.client.models.node import MatterNode
 from matter_server.common.errors import MatterError
@@ -28,6 +27,7 @@ from homeassistant.setup import async_setup_component
 from .common import load_and_parse_node_fixture, setup_integration_with_node_fixture
 
 from tests.common import MockConfigEntry
+from tests.typing import WebSocketGenerator
 
 
 @pytest.fixture(name="connect_timeout")
@@ -613,7 +613,7 @@ async def test_remove_entry(
 async def test_remove_config_entry_device(
     hass: HomeAssistant,
     matter_client: MagicMock,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test that a device can be removed ok."""
     assert await async_setup_component(hass, "config", {})
@@ -656,7 +656,7 @@ async def test_remove_config_entry_device_no_node(
     hass: HomeAssistant,
     matter_client: MagicMock,
     integration: MockConfigEntry,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test that a device can be removed ok without an existing node."""
     assert await async_setup_component(hass, "config", {})

tests/components/mysensors/test_init.py

@@ -1,9 +1,6 @@
 """Test function in __init__.py."""
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
-
-from aiohttp import ClientWebSocketResponse
 from mysensors import BaseSyncGateway
 from mysensors.sensor import Sensor
 
@@ -13,6 +10,7 @@ from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
+from tests.typing import WebSocketGenerator
 
 
 async def test_remove_config_entry_device(
@@ -20,7 +18,7 @@ async def test_remove_config_entry_device(
     gps_sensor: Sensor,
     integration: MockConfigEntry,
     gateway: BaseSyncGateway,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test that a device can be removed ok."""
     entity_id = "sensor.gps_sensor_1_1"

tests/components/onewire/test_init.py

@@ -1,5 +1,4 @@
 """Tests for 1-Wire config flow."""
-from collections.abc import Awaitable, Callable
 from unittest.mock import MagicMock, patch
 
 import aiohttp
@@ -15,6 +14,8 @@ from homeassistant.setup import async_setup_component
 
 from . import setup_owproxy_mock_devices
 
+from tests.typing import WebSocketGenerator
+
 
 async def remove_device(
     ws_client: aiohttp.ClientWebSocketResponse, device_id: str, config_entry_id: str
@@ -78,9 +79,7 @@ async def test_registry_cleanup(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
     owproxy: MagicMock,
-    hass_ws_client: Callable[
+    hass_ws_client: WebSocketGenerator,
-        [HomeAssistant], Awaitable[aiohttp.ClientWebSocketResponse]
-    ],
 ):
     """Test being able to remove a disconnected device."""
     assert await async_setup_component(hass, "config", {})

tests/components/pushover/test_init.py

@@ -1,8 +1,6 @@
 """Test pushbullet integration."""
-from collections.abc import Awaitable, Callable
 from unittest.mock import MagicMock, patch
 
-import aiohttp
 from pushover_complete import BadAPIRequestError
 import pytest
 import requests_mock
@@ -17,6 +15,7 @@ from . import MOCK_CONFIG
 
 from tests.common import MockConfigEntry
 from tests.components.repairs import get_repairs
+from tests.typing import WebSocketGenerator
 
 
 @pytest.fixture(autouse=False)
@@ -30,9 +29,7 @@ def mock_pushover():
 
 async def test_setup(
     hass: HomeAssistant,
-    hass_ws_client: Callable[
+    hass_ws_client: WebSocketGenerator,
-        [HomeAssistant], Awaitable[aiohttp.ClientWebSocketResponse]
-    ],
     mock_pushover: MagicMock,
 ) -> None:
     """Test integration failed due to an error."""

tests/components/repairs/__init__.py

@@ -1,15 +1,15 @@
 """Tests for the repairs integration."""
-from collections.abc import Awaitable, Callable
 
-from aiohttp import ClientWebSocketResponse
 
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
+from tests.typing import WebSocketGenerator
+
 
 async def get_repairs(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ):
     """Return the repairs list of issues."""
     assert await async_setup_component(hass, "repairs", {})

tests/components/repairs/test_init.py

@@ -1,8 +1,6 @@
 """Test the repairs websocket API."""
-from collections.abc import Awaitable, Callable
 from unittest.mock import AsyncMock, Mock
 
-from aiohttp import ClientWebSocketResponse
 from freezegun import freeze_time
 import pytest
 
@@ -489,7 +487,7 @@ async def test_non_compliant_platform(
 @freeze_time("2022-07-21 08:22:00")
 async def test_sync_methods(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test sync method for creating and deleting an issue."""
 

tests/components/repairs/test_websocket_api.py

@@ -1,12 +1,10 @@
 """Test the repairs websocket API."""
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
 from http import HTTPStatus
 from typing import Any
 from unittest.mock import ANY, AsyncMock, Mock
 
-from aiohttp import ClientWebSocketResponse
 from freezegun import freeze_time
 import pytest
 import voluptuous as vol
@@ -524,7 +522,7 @@ async def test_list_issues(
 async def test_fix_issue_aborted(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test we can fix an issue."""
     assert await async_setup_component(hass, "http", {})

tests/components/rtsp_to_webrtc/test_init.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import base64
-from collections.abc import Awaitable, Callable
 from typing import Any
 from unittest.mock import patch
 
@@ -20,6 +19,7 @@ from .conftest import SERVER_URL, STREAM_SOURCE, ComponentSetup
 
 from tests.common import MockConfigEntry
 from tests.test_util.aiohttp import AiohttpClientMocker
+from tests.typing import WebSocketGenerator
 
 # The webrtc component does not inspect the details of the offer and answer,
 # and is only a pass through.
@@ -83,7 +83,7 @@ async def test_setup_communication_failure(
 async def test_offer_for_stream_source(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    hass_ws_client: Callable[[...], Awaitable[aiohttp.ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     mock_camera: Any,
     rtsp_to_webrtc_client: Any,
     setup_integration: ComponentSetup,
@@ -124,7 +124,7 @@ async def test_offer_for_stream_source(
 async def test_offer_failure(
     hass: HomeAssistant,
     aioclient_mock: AiohttpClientMocker,
-    hass_ws_client: Callable[[...], Awaitable[aiohttp.ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     mock_camera: Any,
     rtsp_to_webrtc_client: Any,
     setup_integration: ComponentSetup,
@@ -161,7 +161,7 @@ async def test_no_stun_server(
     hass: HomeAssistant,
     rtsp_to_webrtc_client: Any,
     setup_integration: ComponentSetup,
-    hass_ws_client: Callable[[...], Awaitable[aiohttp.ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test successful setup and unload."""
     await setup_integration()
@@ -188,7 +188,7 @@ async def test_stun_server(
     rtsp_to_webrtc_client: Any,
     setup_integration: ComponentSetup,
     config_entry: MockConfigEntry,
-    hass_ws_client: Callable[[...], Awaitable[aiohttp.ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test successful setup and unload."""
     await setup_integration()

tests/components/schedule/test_init.py

@@ -1,11 +1,10 @@
 """Test for the Schedule integration."""
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable, Coroutine
+from collections.abc import Callable, Coroutine
 from typing import Any
 from unittest.mock import patch
 
-from aiohttp import ClientWebSocketResponse
 import pytest
 
 from homeassistant.components.schedule import STORAGE_VERSION, STORAGE_VERSION_MINOR
@@ -39,6 +38,7 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockUser, async_capture_events, async_fire_time_changed
+from tests.typing import WebSocketGenerator
 
 
 @pytest.fixture
@@ -537,7 +537,7 @@ async def test_schedule_updates(
 
 async def test_ws_list(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     schedule_setup: Callable[..., Coroutine[Any, Any, bool]],
 ) -> None:
     """Test listing via WS."""
@@ -567,7 +567,7 @@ async def test_ws_list(
 
 async def test_ws_delete(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     schedule_setup: Callable[..., Coroutine[Any, Any, bool]],
 ) -> None:
     """Test WS delete cleans up entity registry."""
@@ -602,7 +602,7 @@ async def test_ws_delete(
 )
 async def test_update(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     schedule_setup: Callable[..., Coroutine[Any, Any, bool]],
     to: str,
     next_event: str,
@@ -672,7 +672,7 @@ async def test_update(
 )
 async def test_ws_create(
     hass: HomeAssistant,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
     schedule_setup: Callable[..., Coroutine[Any, Any, bool]],
     freezer,
     to: str,

tests/components/unifiprotect/test_init.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
 from unittest.mock import AsyncMock, patch
 
 import aiohttp
@@ -23,6 +22,7 @@ from . import _patch_discovery
 from .utils import MockUFPFixture, init_entry, time_changed
 
 from tests.common import MockConfigEntry
+from tests.typing import WebSocketGenerator
 
 
 async def remove_device(
@@ -217,9 +217,7 @@ async def test_device_remove_devices(
     hass: HomeAssistant,
     ufp: MockUFPFixture,
     light: Light,
-    hass_ws_client: Callable[
+    hass_ws_client: WebSocketGenerator,
-        [HomeAssistant], Awaitable[aiohttp.ClientWebSocketResponse]
-    ],
 ) -> None:
     """Test we can only remove a device that no longer exists."""
 
@@ -252,9 +250,7 @@ async def test_device_remove_devices(
 async def test_device_remove_devices_nvr(
     hass: HomeAssistant,
     ufp: MockUFPFixture,
-    hass_ws_client: Callable[
+    hass_ws_client: WebSocketGenerator,
-        [HomeAssistant], Awaitable[aiohttp.ClientWebSocketResponse]
-    ],
 ) -> None:
     """Test we can only remove a NVR device that no longer exists."""
     assert await async_setup_component(hass, "config", {})

tests/components/update/test_init.py

@@ -1,8 +1,6 @@
 """The tests for the Update component."""
-from collections.abc import Awaitable, Callable
 from unittest.mock import MagicMock, patch
 
-from aiohttp import ClientWebSocketResponse
 import pytest
 
 from homeassistant.components.update import (
@@ -40,6 +38,7 @@ from homeassistant.helpers.event import async_track_state_change_event
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockEntityPlatform, mock_restore_cache
+from tests.typing import WebSocketGenerator
 
 
 class MockUpdateEntity(UpdateEntity):
@@ -681,7 +680,7 @@ async def test_restore_state(
 async def test_release_notes(
     hass: HomeAssistant,
     enable_custom_integrations: None,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test getting the release notes over the websocket connection."""
     platform = getattr(hass.components, f"test.{DOMAIN}")
@@ -707,7 +706,7 @@ async def test_release_notes(
 async def test_release_notes_entity_not_found(
     hass: HomeAssistant,
     enable_custom_integrations: None,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test getting the release notes for not found entity."""
     platform = getattr(hass.components, f"test.{DOMAIN}")
@@ -734,7 +733,7 @@ async def test_release_notes_entity_not_found(
 async def test_release_notes_entity_does_not_support_release_notes(
     hass: HomeAssistant,
     enable_custom_integrations: None,
-    hass_ws_client: Callable[[HomeAssistant], Awaitable[ClientWebSocketResponse]],
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test getting the release notes for entity that does not support release notes."""
     platform = getattr(hass.components, f"test.{DOMAIN}")