Better handling of balboa spa connection (#71909)

* Better handling of balboa spa connection

* Send a single message for keep alive task rather than multiple

homeassistant/components/balboa/__init__.py

@@ -22,6 +22,7 @@ from .const import (
     SIGNAL_UPDATE,
 )
 
+KEEP_ALIVE_INTERVAL = timedelta(minutes=1)
 SYNC_TIME_INTERVAL = timedelta(days=1)
 
 
@@ -31,7 +32,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     _LOGGER.debug("Attempting to connect to %s", host)
     spa = BalboaSpaWifi(host)
-
     connected = await spa.connect()
     if not connected:
         _LOGGER.error("Failed to connect to spa at %s", host)
@@ -39,11 +39,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = spa
 
-    # send config requests, and then listen until we are configured.
+    async def _async_balboa_update_cb() -> None:
-    await spa.send_mod_ident_req()
-    await spa.send_panel_req(0, 1)
-
-    async def _async_balboa_update_cb():
         """Primary update callback called from pybalboa."""
         _LOGGER.debug("Primary update callback triggered")
         async_dispatcher_send(hass, SIGNAL_UPDATE.format(entry.entry_id))
@@ -52,13 +48,30 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     spa.new_data_cb = _async_balboa_update_cb
 
     _LOGGER.debug("Starting listener and monitor tasks")
-    asyncio.create_task(spa.listen())
+    monitoring_tasks = [asyncio.create_task(spa.listen())]
     await spa.spa_configured()
-    asyncio.create_task(spa.check_connection_status())
+    monitoring_tasks.append(asyncio.create_task(spa.check_connection_status()))
+
+    def stop_monitoring() -> None:
+        """Stop monitoring the spa connection."""
+        _LOGGER.debug("Canceling listener and monitor tasks")
+        for task in monitoring_tasks:
+            task.cancel()
+
+    entry.async_on_unload(stop_monitoring)
 
     # At this point we have a configured spa.
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
+    async def keep_alive(now: datetime) -> None:
+        """Keep alive task."""
+        _LOGGER.debug("Keep alive")
+        await spa.send_mod_ident_req()
+
+    entry.async_on_unload(
+        async_track_time_interval(hass, keep_alive, KEEP_ALIVE_INTERVAL)
+    )
+
     # call update_listener on startup and for options change as well.
     await async_setup_time_sync(hass, entry)
     entry.async_on_unload(entry.add_update_listener(update_listener))
@@ -68,14 +81,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-
     _LOGGER.debug("Disconnecting from spa")
-    spa = hass.data[DOMAIN][entry.entry_id]
+    spa: BalboaSpaWifi = hass.data[DOMAIN][entry.entry_id]
-    await spa.disconnect()
 
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
         hass.data[DOMAIN].pop(entry.entry_id)
 
+    await spa.disconnect()
+
     return unload_ok
 
 
@@ -90,9 +103,9 @@ async def async_setup_time_sync(hass: HomeAssistant, entry: ConfigEntry) -> None
         return
 
     _LOGGER.debug("Setting up daily time sync")
-    spa = hass.data[DOMAIN][entry.entry_id]
+    spa: BalboaSpaWifi = hass.data[DOMAIN][entry.entry_id]
 
-    async def sync_time(now: datetime):
+    async def sync_time(now: datetime) -> None:
         _LOGGER.debug("Syncing time with Home Assistant")
         await spa.set_time(time.strptime(str(dt_util.now()), "%Y-%m-%d %H:%M:%S.%f%z"))
 

homeassistant/components/balboa/config_flow.py

@@ -1,12 +1,16 @@
 """Config flow for Balboa Spa Client integration."""
+from __future__ import annotations
+
 import asyncio
+from typing import Any
 
 from pybalboa import BalboaSpaWifi
 import voluptuous as vol
 
-from homeassistant import config_entries, core, exceptions
+from homeassistant import config_entries, exceptions
 from homeassistant.const import CONF_HOST
 from homeassistant.core import callback
+from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.device_registry import format_mac
 
 from .const import _LOGGER, CONF_SYNC_TIME, DOMAIN
@@ -14,9 +18,8 @@ from .const import _LOGGER, CONF_SYNC_TIME, DOMAIN
 DATA_SCHEMA = vol.Schema({vol.Required(CONF_HOST): str})
 
 
-async def validate_input(hass: core.HomeAssistant, data):
+async def validate_input(data: dict[str, Any]) -> dict[str, str]:
     """Validate the user input allows us to connect."""
-
     _LOGGER.debug("Attempting to connect to %s", data[CONF_HOST])
     spa = BalboaSpaWifi(data[CONF_HOST])
     connected = await spa.connect()
@@ -24,16 +27,12 @@ async def validate_input(hass: core.HomeAssistant, data):
     if not connected:
         raise CannotConnect
 
-    # send config requests, and then listen until we are configured.
+    task = asyncio.create_task(spa.listen())
-    await spa.send_mod_ident_req()
-    await spa.send_panel_req(0, 1)
-
-    asyncio.create_task(spa.listen())
-
     await spa.spa_configured()
 
     mac_addr = format_mac(spa.get_macaddr())
     model = spa.get_model_name()
+    task.cancel()
     await spa.disconnect()
 
     return {"title": model, "formatted_mac": mac_addr}
@@ -46,17 +45,21 @@ class BalboaSpaClientFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
     @staticmethod
     @callback
-    def async_get_options_flow(config_entry):
+    def async_get_options_flow(
+        config_entry: config_entries.ConfigEntry,
+    ) -> config_entries.OptionsFlow:
         """Get the options flow for this handler."""
         return BalboaSpaClientOptionsFlowHandler(config_entry)
 
-    async def async_step_user(self, user_input=None):
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
         """Handle a flow initialized by the user."""
         errors = {}
         if user_input is not None:
             self._async_abort_entries_match({CONF_HOST: user_input[CONF_HOST]})
             try:
-                info = await validate_input(self.hass, user_input)
+                info = await validate_input(user_input)
             except CannotConnect:
                 errors["base"] = "cannot_connect"
             except Exception:  # pylint: disable=broad-except
@@ -79,11 +82,13 @@ class CannotConnect(exceptions.HomeAssistantError):
 class BalboaSpaClientOptionsFlowHandler(config_entries.OptionsFlow):
     """Handle Balboa Spa Client options."""
 
-    def __init__(self, config_entry):
+    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
         """Initialize Balboa Spa Client options flow."""
         self.config_entry = config_entry
 
-    async def async_step_init(self, user_input=None):
+    async def async_step_init(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
         """Manage Balboa Spa Client options."""
         if user_input is not None:
             return self.async_create_entry(title="", data=user_input)