From 7e7a27f8b983e5a803eaacf572b0d779b1dd3999 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 30 Jan 2023 13:40:07 +0100
Subject: [PATCH] Add type hints to http client fixtures (#86795)

---
 tests/conftest.py | 34 +++++++++++++++++++++++++---------
 tests/typing.py   | 11 +++++++++++
 2 files changed, 36 insertions(+), 9 deletions(-)
 create mode 100644 tests/typing.py

diff --git a/tests/conftest.py b/tests/conftest.py
index 6aa4d639d99..5625fcbdfb7 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -16,8 +16,7 @@ import threading
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
-from aiohttp import client
-from aiohttp.pytest_plugin import AiohttpClient
+from aiohttp import ClientWebSocketResponse, client
 from aiohttp.test_utils import (
     BaseTestServer,
     TestClient,
@@ -55,6 +54,7 @@ from homeassistant.setup import async_setup_component
 from homeassistant.util import dt as dt_util, location
 
 from .ignore_uncaught_exceptions import IGNORE_UNCAUGHT_EXCEPTIONS
+from .typing import TestClientGenerator, TestWebSocketGenerator
 
 pytest.register_assert_rewrite("tests.common")
 
@@ -332,7 +332,7 @@ def aiohttp_client_cls():
 @pytest.fixture
 def aiohttp_client(
     event_loop: asyncio.AbstractEventLoop,
-) -> Generator[AiohttpClient, None, None]:
+) -> Generator[TestClientGenerator, None, None]:
     """Override the default aiohttp_client since 3.x does not support aiohttp_client_cls.
 
     Remove this when upgrading to 4.x as aiohttp_client_cls
@@ -607,10 +607,15 @@ def local_auth(hass):
 
 
 @pytest.fixture
-def hass_client(hass, aiohttp_client, hass_access_token, socket_enabled):
+def hass_client(
+    hass: HomeAssistant,
+    aiohttp_client: TestClientGenerator,
+    hass_access_token: str,
+    socket_enabled: None,
+) -> TestClientGenerator:
     """Return an authenticated HTTP client."""
 
-    async def auth_client():
+    async def auth_client() -> TestClient:
         """Return an authenticated client."""
         return await aiohttp_client(
             hass.http.app, headers={"Authorization": f"Bearer {hass_access_token}"}
@@ -620,10 +625,14 @@ def hass_client(hass, aiohttp_client, hass_access_token, socket_enabled):
 
 
 @pytest.fixture
-def hass_client_no_auth(hass, aiohttp_client, socket_enabled):
+def hass_client_no_auth(
+    hass: HomeAssistant,
+    aiohttp_client: TestClientGenerator,
+    socket_enabled: None,
+) -> TestClientGenerator:
     """Return an unauthenticated HTTP client."""
 
-    async def client():
+    async def client() -> TestClient:
         """Return an authenticated client."""
         return await aiohttp_client(hass.http.app)
 
@@ -655,10 +664,17 @@ def current_request_with_host(current_request):
 
 
 @pytest.fixture
-def hass_ws_client(aiohttp_client, hass_access_token, hass, socket_enabled):
+def hass_ws_client(
+    aiohttp_client: TestClientGenerator,
+    hass_access_token: str | None,
+    hass: HomeAssistant,
+    socket_enabled: None,
+) -> TestWebSocketGenerator:
     """Websocket client fixture connected to websocket server."""
 
-    async def create_client(hass=hass, access_token=hass_access_token):
+    async def create_client(
+        hass: HomeAssistant = hass, access_token: str | None = hass_access_token
+    ) -> ClientWebSocketResponse:
         """Create a websocket client."""
         assert await async_setup_component(hass, "websocket_api", {})
         client = await aiohttp_client(hass.http.app)
diff --git a/tests/typing.py b/tests/typing.py
new file mode 100644
index 00000000000..ec6577a9e87
--- /dev/null
+++ b/tests/typing.py
@@ -0,0 +1,11 @@
+"""Typing helpers for Home Assistant tests."""
+from __future__ import annotations
+
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from aiohttp import ClientWebSocketResponse
+from aiohttp.test_utils import TestClient
+
+TestClientGenerator = Callable[..., Coroutine[Any, Any, TestClient]]
+TestWebSocketGenerator = Callable[..., Coroutine[Any, Any, ClientWebSocketResponse]]