Add MockHAClientWebSocket test helper (#87832)

* Add MockHAClientWebSocket test helper

* Add sample use

* Add missing type hint

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"]

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
 

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]]