Patch aiohttp client session close (#34769)

* Patch aiohttp client session close

* Add test

* Restore close regardless of auto_cleanup

* Close session instead of detaching and do not restore

* Delint test

* Add frame helper

* Use frame helper

* Test warning log when closing session

* Clean up

* Correct docstring

* Do not change shutdown

* Fix tests

homeassistant/helpers/aiohttp_client.py

@@ -1,5 +1,6 @@
 """Helper for aiohttp webclient stuff."""
 import asyncio
+import logging
 from ssl import SSLContext
 import sys
 from typing import Any, Awaitable, Optional, Union, cast
@@ -12,10 +13,13 @@ import async_timeout
 
 from homeassistant.const import EVENT_HOMEASSISTANT_CLOSE, __version__
 from homeassistant.core import Event, callback
+from homeassistant.helpers.frame import MissingIntegrationFrame, get_integration_frame
 from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.loader import bind_hass
 from homeassistant.util import ssl as ssl_util
 
+_LOGGER = logging.getLogger(__name__)
+
 DATA_CONNECTOR = "aiohttp_connector"
 DATA_CONNECTOR_NOTVERIFY = "aiohttp_connector_notverify"
 DATA_CLIENTSESSION = "aiohttp_clientsession"
@@ -67,6 +71,35 @@ def async_create_clientsession(
         connector=connector, headers={USER_AGENT: SERVER_SOFTWARE}, **kwargs,
     )
 
+    async def patched_close() -> None:
+        """Mock close to avoid integrations closing our session."""
+        try:
+            found_frame, integration, path = get_integration_frame()
+        except MissingIntegrationFrame:
+            # Did not source from an integration? Hard error.
+            raise RuntimeError(
+                "Detected closing of the Home Assistant aiohttp session in the Home Assistant core. "
+                "Please report this issue."
+            )
+
+        index = found_frame.filename.index(path)
+        if path == "custom_components/":
+            extra = " to the custom component author"
+        else:
+            extra = ""
+
+        _LOGGER.warning(
+            "Detected integration that closes the Home Assistant aiohttp session. "
+            "Please report issue%s for %s using this method at %s, line %s: %s",
+            extra,
+            integration,
+            found_frame.filename[index:],
+            found_frame.lineno,
+            found_frame.line.strip(),
+        )
+
+    clientsession.close = patched_close  # type: ignore
+
     if auto_cleanup:
         _async_register_clientsession_shutdown(hass, clientsession)
 

homeassistant/helpers/frame.py

@@ -0,0 +1,36 @@
+"""Provide frame helper for finding the current frame context."""
+from traceback import FrameSummary, extract_stack
+from typing import Tuple
+
+from homeassistant.exceptions import HomeAssistantError
+
+
+def get_integration_frame() -> Tuple[FrameSummary, str, str]:
+    """Return the frame, integration and integration path of the current stack frame."""
+    found_frame = None
+
+    for frame in reversed(extract_stack()):
+        for path in ("custom_components/", "homeassistant/components/"):
+            try:
+                index = frame.filename.index(path)
+                found_frame = frame
+                break
+            except ValueError:
+                continue
+
+        if found_frame is not None:
+            break
+
+    if found_frame is None:
+        raise MissingIntegrationFrame
+
+    start = index + len(path)
+    end = found_frame.filename.index("/", start)
+
+    integration = found_frame.filename[start:end]
+
+    return found_frame, integration, path
+
+
+class MissingIntegrationFrame(HomeAssistantError):
+    """Raised when no integration is found in the frame."""

tests/helpers/test_aiohttp_client.py

@@ -8,9 +8,11 @@ from homeassistant.core import EVENT_HOMEASSISTANT_CLOSE
 import homeassistant.helpers.aiohttp_client as client
 from homeassistant.setup import async_setup_component
 
+from tests.async_mock import Mock, patch
 
-@pytest.fixture
-def camera_client(hass, hass_client):
+
+@pytest.fixture(name="camera_client")
+def camera_client_fixture(hass, hass_client):
     """Fixture to fetch camera streams."""
     assert hass.loop.run_until_complete(
         async_setup_component(
@@ -91,6 +93,82 @@ async def test_get_clientsession_cleanup_without_ssl(hass):
     assert hass.data[client.DATA_CONNECTOR_NOTVERIFY].closed
 
 
+async def test_get_clientsession_patched_close(hass):
+    """Test closing clientsession does not work."""
+    with patch("aiohttp.ClientSession.close") as mock_close:
+        session = client.async_get_clientsession(hass)
+
+        assert isinstance(hass.data[client.DATA_CLIENTSESSION], aiohttp.ClientSession)
+        assert isinstance(hass.data[client.DATA_CONNECTOR], aiohttp.TCPConnector)
+
+        with pytest.raises(RuntimeError):
+            await session.close()
+
+        assert mock_close.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.close()",
+            ),
+            Mock(
+                filename="/home/paulus/aiohue/lights.py",
+                lineno="2",
+                line="something()",
+            ),
+        ],
+    ):
+        session = client.async_get_clientsession(hass)
+        await session.close()
+    assert (
+        "Detected integration that closes the Home Assistant aiohttp session. "
+        "Please report issue for hue using this method at "
+        "homeassistant/components/hue/light.py, line 23: await session.close()"
+    ) 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.close()",
+            ),
+            Mock(
+                filename="/home/paulus/aiohue/lights.py",
+                lineno="2",
+                line="something()",
+            ),
+        ],
+    ):
+        session = client.async_get_clientsession(hass)
+        await session.close()
+    assert (
+        "Detected integration that closes the Home Assistant aiohttp session. "
+        "Please report issue to the custom component author for hue using this method at "
+        "custom_components/hue/light.py, line 23: await session.close()" in caplog.text
+    )
+
+
 async def test_async_aiohttp_proxy_stream(aioclient_mock, camera_client):
     """Test that it fetches the given url."""
     aioclient_mock.get("http://example.com/mjpeg_stream", content=b"Frame1Frame2Frame3")

tests/helpers/test_frame.py

@@ -0,0 +1,56 @@
+"""Test the frame helper."""
+import pytest
+
+from homeassistant.helpers import frame
+
+from tests.async_mock import Mock, patch
+
+
+async def test_extract_frame_integration(caplog):
+    """Test extracting the current frame from integration context."""
+    correct_frame = Mock(
+        filename="/home/paulus/homeassistant/components/hue/light.py",
+        lineno="23",
+        line="self.light.is_on",
+    )
+    with patch(
+        "homeassistant.helpers.frame.extract_stack",
+        return_value=[
+            Mock(
+                filename="/home/paulus/homeassistant/core.py",
+                lineno="23",
+                line="do_something()",
+            ),
+            correct_frame,
+            Mock(
+                filename="/home/paulus/aiohue/lights.py",
+                lineno="2",
+                line="something()",
+            ),
+        ],
+    ):
+        found_frame, integration, path = frame.get_integration_frame()
+
+    assert integration == "hue"
+    assert path == "homeassistant/components/"
+    assert found_frame == correct_frame
+
+
+async def test_extract_frame_no_integration(caplog):
+    """Test extracting the current frame without 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/aiohue/lights.py",
+                lineno="2",
+                line="something()",
+            ),
+        ],
+    ), pytest.raises(frame.MissingIntegrationFrame):
+        frame.get_integration_frame()