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