Create httpx helper to wrap a shared httpx.AsyncClient (#43877)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

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()

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:

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:

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

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()

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)

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):

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
+    )