From dd0afc3b66b8e3826e31a72ac7566d9e62481c41 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 9 Dec 2020 10:18:57 -0600 Subject: [PATCH] Create httpx helper to wrap a shared httpx.AsyncClient (#43877) Co-authored-by: Paulus Schoutsen --- homeassistant/components/pvoutput/sensor.py | 20 ++- .../components/rest/binary_sensor.py | 8 +- homeassistant/components/rest/data.py | 13 +- homeassistant/components/rest/sensor.py | 15 +- homeassistant/components/scrape/sensor.py | 14 +- homeassistant/helpers/httpx_client.py | 88 +++++++++++ tests/components/rest/test_sensor.py | 13 ++ tests/helpers/test_httpx_client.py | 143 ++++++++++++++++++ 8 files changed, 284 insertions(+), 30 deletions(-) create mode 100644 homeassistant/helpers/httpx_client.py create mode 100644 tests/helpers/test_httpx_client.py diff --git a/homeassistant/components/pvoutput/sensor.py b/homeassistant/components/pvoutput/sensor.py index fb3446fb652..32d33f19e80 100644 --- a/homeassistant/components/pvoutput/sensor.py +++ b/homeassistant/components/pvoutput/sensor.py @@ -15,6 +15,7 @@ from homeassistant.const import ( CONF_API_KEY, CONF_NAME, ) +from homeassistant.core import callback import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity import Entity @@ -53,14 +54,14 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info= verify_ssl = DEFAULT_VERIFY_SSL headers = {"X-Pvoutput-Apikey": api_key, "X-Pvoutput-SystemId": system_id} - rest = RestData(method, _ENDPOINT, auth, headers, None, payload, verify_ssl) + rest = RestData(hass, method, _ENDPOINT, auth, headers, None, payload, verify_ssl) await rest.async_update() if rest.data is None: _LOGGER.error("Unable to fetch data from PVOutput") return False - async_add_entities([PvoutputSensor(rest, name)], True) + async_add_entities([PvoutputSensor(rest, name)]) class PvoutputSensor(Entity): @@ -114,13 +115,18 @@ class PvoutputSensor(Entity): async def async_update(self): """Get the latest data from the PVOutput API and updates the state.""" + await self.rest.async_update() + self._async_update_from_rest_data() + + async def async_added_to_hass(self): + """Ensure the data from the initial update is reflected in the state.""" + self._async_update_from_rest_data() + + @callback + def _async_update_from_rest_data(self): + """Update state from the rest data.""" try: - await self.rest.async_update() self.pvcoutput = self.status._make(self.rest.data.split(",")) except TypeError: self.pvcoutput = None _LOGGER.error("Unable to fetch data from PVOutput. %s", self.rest.data) - - async def async_will_remove_from_hass(self): - """Shutdown the session.""" - await self.rest.async_remove() diff --git a/homeassistant/components/rest/binary_sensor.py b/homeassistant/components/rest/binary_sensor.py index 7f0f920b843..49c10354c51 100644 --- a/homeassistant/components/rest/binary_sensor.py +++ b/homeassistant/components/rest/binary_sensor.py @@ -101,9 +101,10 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info= auth = None rest = RestData( - method, resource, auth, headers, params, payload, verify_ssl, timeout + hass, method, resource, auth, headers, params, payload, verify_ssl, timeout ) await rest.async_update() + if rest.data is None: raise PlatformNotReady @@ -119,7 +120,6 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info= resource_template, ) ], - True, ) @@ -187,10 +187,6 @@ class RestBinarySensor(BinarySensorEntity): """Force update.""" return self._force_update - async def async_will_remove_from_hass(self): - """Shutdown the session.""" - await self.rest.async_remove() - async def async_update(self): """Get the latest data from REST API and updates the state.""" if self._resource_template is not None: diff --git a/homeassistant/components/rest/data.py b/homeassistant/components/rest/data.py index bd35383e981..dd2e29616c7 100644 --- a/homeassistant/components/rest/data.py +++ b/homeassistant/components/rest/data.py @@ -3,6 +3,8 @@ import logging import httpx +from homeassistant.helpers.httpx_client import get_async_client + DEFAULT_TIMEOUT = 10 _LOGGER = logging.getLogger(__name__) @@ -13,6 +15,7 @@ class RestData: def __init__( self, + hass, method, resource, auth, @@ -23,6 +26,7 @@ class RestData: timeout=DEFAULT_TIMEOUT, ): """Initialize the data object.""" + self._hass = hass self._method = method self._resource = resource self._auth = auth @@ -35,11 +39,6 @@ class RestData: self.data = None self.headers = None - async def async_remove(self): - """Destroy the http session on destroy.""" - if self._async_client: - await self._async_client.aclose() - def set_url(self, url): """Set url.""" self._resource = url @@ -47,7 +46,9 @@ class RestData: async def async_update(self): """Get the latest data from REST service with provided method.""" if not self._async_client: - self._async_client = httpx.AsyncClient(verify=self._verify_ssl) + self._async_client = get_async_client( + self._hass, verify_ssl=self._verify_ssl + ) _LOGGER.debug("Updating from %s", self._resource) try: diff --git a/homeassistant/components/rest/sensor.py b/homeassistant/components/rest/sensor.py index f048eaa3b47..51d9ec20472 100644 --- a/homeassistant/components/rest/sensor.py +++ b/homeassistant/components/rest/sensor.py @@ -116,8 +116,9 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info= else: auth = None rest = RestData( - method, resource, auth, headers, params, payload, verify_ssl, timeout + hass, method, resource, auth, headers, params, payload, verify_ssl, timeout ) + await rest.async_update() if rest.data is None: @@ -140,7 +141,6 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info= json_attrs_path, ) ], - True, ) @@ -210,7 +210,14 @@ class RestSensor(Entity): self.rest.set_url(self._resource_template.async_render(parse_result=False)) await self.rest.async_update() + self._update_from_rest_data() + async def async_added_to_hass(self): + """Ensure the data from the initial update is reflected in the state.""" + self._update_from_rest_data() + + def _update_from_rest_data(self): + """Update state from the rest data.""" value = self.rest.data _LOGGER.debug("Data fetched from resource: %s", value) if self.rest.headers is not None: @@ -266,10 +273,6 @@ class RestSensor(Entity): self._state = value - async def async_will_remove_from_hass(self): - """Shutdown the session.""" - await self.rest.async_remove() - @property def device_state_attributes(self): """Return the state attributes.""" diff --git a/homeassistant/components/scrape/sensor.py b/homeassistant/components/scrape/sensor.py index b76995fe39f..e8b6fcfd2c3 100644 --- a/homeassistant/components/scrape/sensor.py +++ b/homeassistant/components/scrape/sensor.py @@ -78,7 +78,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info= auth = HTTPBasicAuth(username, password) else: auth = None - rest = RestData(method, resource, auth, headers, None, payload, verify_ssl) + rest = RestData(hass, method, resource, auth, headers, None, payload, verify_ssl) await rest.async_update() if rest.data is None: @@ -137,6 +137,14 @@ class ScrapeSensor(Entity): async def async_update(self): """Get the latest data from the source and updates the state.""" await self.rest.async_update() + await self._async_update_from_rest_data() + + async def async_added_to_hass(self): + """Ensure the data from the initial update is reflected in the state.""" + await self._async_update_from_rest_data() + + async def _async_update_from_rest_data(self): + """Update state from the rest data.""" if self.rest.data is None: _LOGGER.error("Unable to retrieve data for %s", self.name) return @@ -153,7 +161,3 @@ class ScrapeSensor(Entity): ) else: self._state = value - - async def async_will_remove_from_hass(self): - """Shutdown the session.""" - await self.rest.async_remove() diff --git a/homeassistant/helpers/httpx_client.py b/homeassistant/helpers/httpx_client.py new file mode 100644 index 00000000000..0f1719b388d --- /dev/null +++ b/homeassistant/helpers/httpx_client.py @@ -0,0 +1,88 @@ +"""Helper for httpx.""" +import sys +from typing import Any, Callable, Optional + +import httpx + +from homeassistant.const import EVENT_HOMEASSISTANT_CLOSE, __version__ +from homeassistant.core import Event, callback +from homeassistant.helpers.frame import warn_use +from homeassistant.helpers.typing import HomeAssistantType +from homeassistant.loader import bind_hass + +DATA_ASYNC_CLIENT = "httpx_async_client" +DATA_ASYNC_CLIENT_NOVERIFY = "httpx_async_client_noverify" +SERVER_SOFTWARE = "HomeAssistant/{0} httpx/{1} Python/{2[0]}.{2[1]}".format( + __version__, httpx.__version__, sys.version_info +) +USER_AGENT = "User-Agent" + + +@callback +@bind_hass +def get_async_client( + hass: HomeAssistantType, verify_ssl: bool = True +) -> httpx.AsyncClient: + """Return default httpx AsyncClient. + + This method must be run in the event loop. + """ + key = DATA_ASYNC_CLIENT if verify_ssl else DATA_ASYNC_CLIENT_NOVERIFY + + client: Optional[httpx.AsyncClient] = hass.data.get(key) + + if client is None: + client = hass.data[key] = create_async_httpx_client(hass, verify_ssl) + + return client + + +@callback +def create_async_httpx_client( + hass: HomeAssistantType, + verify_ssl: bool = True, + auto_cleanup: bool = True, + **kwargs: Any, +) -> httpx.AsyncClient: + """Create a new httpx.AsyncClient with kwargs, i.e. for cookies. + + If auto_cleanup is False, the client will be + automatically closed on homeassistant_stop. + + This method must be run in the event loop. + """ + + client = httpx.AsyncClient( + verify=verify_ssl, + headers={USER_AGENT: SERVER_SOFTWARE}, + **kwargs, + ) + + original_aclose = client.aclose + + client.aclose = warn_use( # type: ignore + client.aclose, "closes the Home Assistant httpx client" + ) + + if auto_cleanup: + _async_register_async_client_shutdown(hass, client, original_aclose) + + return client + + +@callback +def _async_register_async_client_shutdown( + hass: HomeAssistantType, + client: httpx.AsyncClient, + original_aclose: Callable[..., Any], +) -> None: + """Register httpx AsyncClient aclose on Home Assistant shutdown. + + This method must be run in the event loop. + """ + + async def _async_close_client(event: Event) -> None: + """Close httpx client.""" + await original_aclose() + + hass.bus.async_listen_once(EVENT_HOMEASSISTANT_CLOSE, _async_close_client) diff --git a/tests/components/rest/test_sensor.py b/tests/components/rest/test_sensor.py index f378a1fc2a1..16d3f8ba0ac 100644 --- a/tests/components/rest/test_sensor.py +++ b/tests/components/rest/test_sensor.py @@ -6,8 +6,10 @@ import httpx import respx from homeassistant import config as hass_config +from homeassistant.components.homeassistant import SERVICE_UPDATE_ENTITY import homeassistant.components.sensor as sensor from homeassistant.const import ( + ATTR_ENTITY_ID, ATTR_UNIT_OF_MEASUREMENT, CONTENT_TYPE_JSON, DATA_MEGABYTES, @@ -151,10 +153,21 @@ async def test_setup_get(hass): } }, ) + await async_setup_component(hass, "homeassistant", {}) await hass.async_block_till_done() assert len(hass.states.async_all()) == 1 + assert hass.states.get("sensor.foo").state == "" + await hass.services.async_call( + "homeassistant", + SERVICE_UPDATE_ENTITY, + {ATTR_ENTITY_ID: "sensor.foo"}, + blocking=True, + ) + await hass.async_block_till_done() + assert hass.states.get("sensor.foo").state == "" + @respx.mock async def test_setup_get_digest_auth(hass): diff --git a/tests/helpers/test_httpx_client.py b/tests/helpers/test_httpx_client.py new file mode 100644 index 00000000000..5444cd4643d --- /dev/null +++ b/tests/helpers/test_httpx_client.py @@ -0,0 +1,143 @@ +"""Test the httpx client helper.""" + +import httpx +import pytest + +from homeassistant.core import EVENT_HOMEASSISTANT_CLOSE +import homeassistant.helpers.httpx_client as client + +from tests.async_mock import Mock, patch + + +async def test_get_async_client_with_ssl(hass): + """Test init async client with ssl.""" + client.get_async_client(hass) + + assert isinstance(hass.data[client.DATA_ASYNC_CLIENT], httpx.AsyncClient) + + +async def test_get_async_client_without_ssl(hass): + """Test init async client without ssl.""" + client.get_async_client(hass, verify_ssl=False) + + assert isinstance(hass.data[client.DATA_ASYNC_CLIENT_NOVERIFY], httpx.AsyncClient) + + +async def test_create_async_httpx_client_with_ssl_and_cookies(hass): + """Test init async client with ssl and cookies.""" + client.get_async_client(hass) + + httpx_client = client.create_async_httpx_client(hass, cookies={"bla": True}) + assert isinstance(httpx_client, httpx.AsyncClient) + assert hass.data[client.DATA_ASYNC_CLIENT] != httpx_client + + +async def test_create_async_httpx_client_without_ssl_and_cookies(hass): + """Test init async client without ssl and cookies.""" + client.get_async_client(hass, verify_ssl=False) + + httpx_client = client.create_async_httpx_client( + hass, verify_ssl=False, cookies={"bla": True} + ) + assert isinstance(httpx_client, httpx.AsyncClient) + assert hass.data[client.DATA_ASYNC_CLIENT_NOVERIFY] != httpx_client + + +async def test_get_async_client_cleanup(hass): + """Test init async client with ssl.""" + client.get_async_client(hass) + + assert isinstance(hass.data[client.DATA_ASYNC_CLIENT], httpx.AsyncClient) + + hass.bus.async_fire(EVENT_HOMEASSISTANT_CLOSE) + await hass.async_block_till_done() + + assert hass.data[client.DATA_ASYNC_CLIENT].is_closed + + +async def test_get_async_client_cleanup_without_ssl(hass): + """Test init async client without ssl.""" + client.get_async_client(hass, verify_ssl=False) + + assert isinstance(hass.data[client.DATA_ASYNC_CLIENT_NOVERIFY], httpx.AsyncClient) + + hass.bus.async_fire(EVENT_HOMEASSISTANT_CLOSE) + await hass.async_block_till_done() + + assert hass.data[client.DATA_ASYNC_CLIENT_NOVERIFY].is_closed + + +async def test_get_async_client_patched_close(hass): + """Test closing the async client does not work.""" + + with patch("httpx.AsyncClient.aclose") as mock_aclose: + httpx_session = client.get_async_client(hass) + assert isinstance(hass.data[client.DATA_ASYNC_CLIENT], httpx.AsyncClient) + + with pytest.raises(RuntimeError): + await httpx_session.aclose() + + assert mock_aclose.call_count == 0 + + +async def test_warning_close_session_integration(hass, caplog): + """Test log warning message when closing the session from integration context.""" + with patch( + "homeassistant.helpers.frame.extract_stack", + return_value=[ + Mock( + filename="/home/paulus/homeassistant/core.py", + lineno="23", + line="do_something()", + ), + Mock( + filename="/home/paulus/homeassistant/components/hue/light.py", + lineno="23", + line="await session.aclose()", + ), + Mock( + filename="/home/paulus/aiohue/lights.py", + lineno="2", + line="something()", + ), + ], + ): + httpx_session = client.get_async_client(hass) + await httpx_session.aclose() + + assert ( + "Detected integration that closes the Home Assistant httpx client. " + "Please report issue for hue using this method at " + "homeassistant/components/hue/light.py, line 23: await session.aclose()" + ) in caplog.text + + +async def test_warning_close_session_custom(hass, caplog): + """Test log warning message when closing the session from custom context.""" + with patch( + "homeassistant.helpers.frame.extract_stack", + return_value=[ + Mock( + filename="/home/paulus/homeassistant/core.py", + lineno="23", + line="do_something()", + ), + Mock( + filename="/home/paulus/config/custom_components/hue/light.py", + lineno="23", + line="await session.aclose()", + ), + Mock( + filename="/home/paulus/aiohue/lights.py", + lineno="2", + line="something()", + ), + ], + ): + httpx_session = client.get_async_client(hass) + await httpx_session.aclose() + assert ( + "Detected integration that closes the Home Assistant httpx client. " + "Please report issue to the custom component author for hue using this method at " + "custom_components/hue/light.py, line 23: await session.aclose()" in caplog.text + )