Skip to content

Commit

Permalink
Limit rainbird aiohttp client session to a single connection (#112146)
Browse files Browse the repository at this point in the history
Limit rainbird to a single open http connection
  • Loading branch information
allenporter authored and balloob committed Mar 4, 2024
1 parent 93ee900 commit 274ab23
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 6 deletions.
7 changes: 4 additions & 3 deletions homeassistant/components/rainbird/__init__.py
Expand Up @@ -11,11 +11,10 @@
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.device_registry import format_mac

from .const import CONF_SERIAL_NUMBER
from .coordinator import RainbirdData
from .coordinator import RainbirdData, async_create_clientsession

_LOGGER = logging.getLogger(__name__)

Expand All @@ -36,9 +35,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:

hass.data.setdefault(DOMAIN, {})

clientsession = async_create_clientsession()
entry.async_on_unload(clientsession.close)
controller = AsyncRainbirdController(
AsyncRainbirdClient(
async_get_clientsession(hass),
clientsession,
entry.data[CONF_HOST],
entry.data[CONF_PASSWORD],
)
Expand Down
7 changes: 5 additions & 2 deletions homeassistant/components/rainbird/config_flow.py
Expand Up @@ -20,7 +20,6 @@
from homeassistant.core import callback
from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers import config_validation as cv, selector
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.device_registry import format_mac

from .const import (
Expand All @@ -30,6 +29,7 @@
DOMAIN,
TIMEOUT_SECONDS,
)
from .coordinator import async_create_clientsession

_LOGGER = logging.getLogger(__name__)

Expand Down Expand Up @@ -101,9 +101,10 @@ async def _test_connection(
Raises a ConfigFlowError on failure.
"""
clientsession = async_create_clientsession()
controller = AsyncRainbirdController(
AsyncRainbirdClient(
async_get_clientsession(self.hass),
clientsession,
host,
password,
)
Expand All @@ -124,6 +125,8 @@ async def _test_connection(
f"Error connecting to Rain Bird controller: {str(err)}",
"cannot_connect",
) from err
finally:
await clientsession.close()

async def async_finish(
self,
Expand Down
11 changes: 11 additions & 0 deletions homeassistant/components/rainbird/coordinator.py
Expand Up @@ -9,6 +9,7 @@
import logging
from typing import TypeVar

import aiohttp
from pyrainbird.async_client import (
AsyncRainbirdController,
RainbirdApiException,
Expand All @@ -28,6 +29,9 @@
# changes, so we refresh it less often.
CALENDAR_UPDATE_INTERVAL = datetime.timedelta(minutes=15)

# Rainbird devices can only accept a single request at a time
CONECTION_LIMIT = 1

_LOGGER = logging.getLogger(__name__)

_T = TypeVar("_T")
Expand All @@ -43,6 +47,13 @@ class RainbirdDeviceState:
rain_delay: int


def async_create_clientsession() -> aiohttp.ClientSession:
"""Create a rainbird async_create_clientsession with a connection limit."""
return aiohttp.ClientSession(
connector=aiohttp.TCPConnector(limit=CONECTION_LIMIT),
)


class RainbirdUpdateCoordinator(DataUpdateCoordinator[RainbirdDeviceState]):
"""Coordinator for rainbird API calls."""

Expand Down
28 changes: 27 additions & 1 deletion tests/components/rainbird/conftest.py
Expand Up @@ -2,6 +2,7 @@

from __future__ import annotations

from collections.abc import Generator
from http import HTTPStatus
import json
from typing import Any
Expand All @@ -15,7 +16,7 @@
ATTR_DURATION,
DEFAULT_TRIGGER_TIME_MINUTES,
)
from homeassistant.const import Platform
from homeassistant.const import EVENT_HOMEASSISTANT_CLOSE, Platform
from homeassistant.core import HomeAssistant

from tests.common import MockConfigEntry
Expand Down Expand Up @@ -155,6 +156,31 @@ def setup_platforms(
yield


@pytest.fixture(autouse=True)
def aioclient_mock(hass: HomeAssistant) -> Generator[AiohttpClientMocker, None, None]:
"""Context manager to mock aiohttp client."""
mocker = AiohttpClientMocker()

def create_session():
session = mocker.create_session(hass.loop)

async def close_session(event):
"""Close session."""
await session.close()

hass.bus.async_listen_once(EVENT_HOMEASSISTANT_CLOSE, close_session)
return session

with patch(
"homeassistant.components.rainbird.async_create_clientsession",
side_effect=create_session,
), patch(
"homeassistant.components.rainbird.config_flow.async_create_clientsession",
side_effect=create_session,
):
yield mocker


def rainbird_json_response(result: dict[str, str]) -> bytes:
"""Create a fake API response."""
return encryption.encrypt(
Expand Down

0 comments on commit 274ab23

Please sign in to comment.