diff --git a/tests/components/analytics/test_init.py b/tests/components/analytics/test_init.py index 0ce4dad5702..a9cb5d28767 100644 --- a/tests/components/analytics/test_init.py +++ b/tests/components/analytics/test_init.py @@ -30,20 +30,20 @@ async def test_websocket( await hass.async_block_till_done() ws_client = await hass_ws_client(hass) - await ws_client.send_json({"id": 1, "type": "analytics"}) + await ws_client.send_json_auto_id({"type": "analytics"}) response = await ws_client.receive_json() assert response["success"] with patch("homeassistant.components.analytics.analytics.HA_VERSION", MOCK_VERSION): - await ws_client.send_json( - {"id": 2, "type": "analytics/preferences", "preferences": {"base": True}} + await ws_client.send_json_auto_id( + {"type": "analytics/preferences", "preferences": {"base": True}} ) response = await ws_client.receive_json() assert len(aioclient_mock.mock_calls) == 1 assert response["result"]["preferences"]["base"] - await ws_client.send_json({"id": 3, "type": "analytics"}) + await ws_client.send_json_auto_id({"type": "analytics"}) response = await ws_client.receive_json() assert response["result"]["preferences"]["base"] diff --git a/tests/conftest.py b/tests/conftest.py index 1f87d2ea351..7628d50a412 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,7 @@ import threading from typing import TYPE_CHECKING, Any, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch -from aiohttp import ClientWebSocketResponse, client +from aiohttp import client from aiohttp.test_utils import ( BaseTestServer, TestClient, @@ -62,6 +62,7 @@ from homeassistant.util import dt as dt_util, location from .ignore_uncaught_exceptions import IGNORE_UNCAUGHT_EXCEPTIONS from .typing import ( ClientSessionGenerator, + MockHAClientWebSocket, MqttMockHAClient, MqttMockHAClientGenerator, MqttMockPahoClient, @@ -725,7 +726,7 @@ def hass_ws_client( async def create_client( hass: HomeAssistant = hass, access_token: str | None = hass_access_token - ) -> ClientWebSocketResponse: + ) -> MockHAClientWebSocket: """Create a websocket client.""" assert await async_setup_component(hass, "websocket_api", {}) client = await aiohttp_client(hass.http.app) @@ -741,9 +742,22 @@ def hass_ws_client( auth_ok = await websocket.receive_json() assert auth_ok["type"] == TYPE_AUTH_OK + def _get_next_id() -> Generator[int, None, None]: + i = 0 + while True: + yield (i := i + 1) + + id_generator = _get_next_id() + + def _send_json_auto_id(data: dict[str, Any]) -> Coroutine[Any, Any, None]: + data["id"] = next(id_generator) + return websocket.send_json(data) + # wrap in client - websocket.client = client - return websocket + wrapped_websocket = cast(MockHAClientWebSocket, websocket) + wrapped_websocket.client = client + wrapped_websocket.send_json_auto_id = _send_json_auto_id + return wrapped_websocket return create_client diff --git a/tests/typing.py b/tests/typing.py index 279fdaaeb2a..7c5391d132c 100644 --- a/tests/typing.py +++ b/tests/typing.py @@ -13,6 +13,14 @@ if TYPE_CHECKING: # testcase which does not use the recorder. from homeassistant.components.recorder import Recorder + +class MockHAClientWebSocket(ClientWebSocketResponse): + """Protocol for a wrapped ClientWebSocketResponse.""" + + client: TestClient + send_json_auto_id: Callable[[dict[str, Any]], Coroutine[Any, Any, None]] + + ClientSessionGenerator = Callable[..., Coroutine[Any, Any, TestClient]] MqttMockPahoClient = MagicMock """MagicMock for `paho.mqtt.client.Client`""" @@ -22,4 +30,4 @@ MqttMockHAClientGenerator = Callable[..., Coroutine[Any, Any, MqttMockHAClient]] """MagicMock generator for `homeassistant.components.mqtt.MQTT`.""" RecorderInstanceGenerator: TypeAlias = Callable[..., Coroutine[Any, Any, "Recorder"]] """Instance generator for `homeassistant.components.recorder.Recorder`.""" -WebSocketGenerator = Callable[..., Coroutine[Any, Any, ClientWebSocketResponse]] +WebSocketGenerator = Callable[..., Coroutine[Any, Any, MockHAClientWebSocket]]