Let config flows temporarily hold a Modbus unit (#179939)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

homeassistant/components/modbus/__init__.py

@@ -9,7 +9,7 @@ from homeassistant.helpers.reload import async_integration_yaml_config
 from homeassistant.helpers.service import async_register_admin_service
 from homeassistant.helpers.typing import ConfigType
 
-from .connection import async_get_unit
+from .connection import async_get_temporary_unit, async_get_unit
 from .const import DOMAIN
 from .modbus import DATA_MODBUS_HUBS, ModbusHub, async_modbus_setup
 from .schemas import CONFIG_SCHEMA
@@ -17,6 +17,7 @@ from .schemas import CONFIG_SCHEMA
 __all__ = [
     "CONFIG_SCHEMA",
     "ModbusHub",
+    "async_get_temporary_unit",
     "async_get_unit",
     "get_hub",
 ]

homeassistant/components/modbus/connection.py

@@ -1,7 +1,10 @@
 """Hand out Modbus units over connections shared between integrations."""
 
+from collections.abc import AsyncIterator, Callable, Coroutine
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 import logging
+from typing import Any
 
 from modbus_connection import (
     ModbusSerialParams,
@@ -41,17 +44,10 @@ class _SharedConnection:
 
 
 @callback
-def async_get_unit(
-    hass: HomeAssistant,
-    entry: ConfigEntry,
-    params: ModbusParams,
-    unit_id: int,
-) -> ModbusUnit:
-    """Return a unit on the connection these credentials describe.
-
-    Consumers of one device share a connection, so their requests serialize
-    behind its lock. It is closed when the last config entry holding a unit on
-    it unloads.
+def _async_acquire(
+    hass: HomeAssistant, params: ModbusParams
+) -> tuple[ModbusConnection, Callable[[], Coroutine[Any, Any, None]]]:
+    """Take a hold on the connection these credentials describe.
 
     Raises `HomeAssistantError` if the device is already in use over different
     link settings, which cannot both be honoured on one connection.
@@ -70,7 +66,7 @@ def async_get_unit(
 
     shared.consumers += 1
 
-    async def _release() -> None:
+    async def release() -> None:
         """Give up this hold, closing behind the last one."""
         shared.consumers -= 1
         if shared.consumers or connections.get(endpoint) is not shared:
@@ -79,5 +75,47 @@ def async_get_unit(
         _LOGGER.debug("Closing the Modbus connection to %s", endpoint)
         await shared.connection.close()
 
-    entry.async_on_unload(_release)
-    return shared.connection.for_unit(unit_id)
+    return shared.connection, release
+
+
+@callback
+def async_get_unit(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    params: ModbusParams,
+    unit_id: int,
+) -> ModbusUnit:
+    """Return a unit on the connection these credentials describe.
+
+    Consumers of one device share a connection, so their requests serialize
+    behind its lock. It is closed when the last config entry holding a unit on
+    it unloads.
+
+    Raises `HomeAssistantError` if the device is already in use over different
+    link settings, which cannot both be honoured on one connection.
+    """
+    connection, release = _async_acquire(hass, params)
+    entry.async_on_unload(release)
+    return connection.for_unit(unit_id)
+
+
+@asynccontextmanager
+async def async_get_temporary_unit(
+    hass: HomeAssistant,
+    params: ModbusParams,
+    unit_id: int,
+) -> AsyncIterator[ModbusUnit]:
+    """Hold a unit on the connection these credentials describe for the context.
+
+    For config flows, which have no config entry yet to tie a hold to. A
+    connection already held by a config entry is shared and stays up; one
+    opened here is closed on exit.
+
+    Raises `HomeAssistantError` if the device is already in use over different
+    link settings, which cannot both be honoured on one connection.
+    """
+    connection, release = _async_acquire(hass, params)
+    try:
+        yield connection.for_unit(unit_id)
+    finally:
+        await release()

tests/components/modbus/test_connection.py

@@ -4,10 +4,12 @@ from collections.abc import Callable, Generator
 from unittest.mock import AsyncMock, patch
 
 from modbus_connection import ModbusSerialParams, ModbusTcpParams
+from modbus_connection.tmodbus import ModbusConnection
 import pytest
 
 from homeassistant.components.modbus.connection import (
     DATA_MODBUS_CONNECTIONS,
+    async_get_temporary_unit,
     async_get_unit,
 )
 from homeassistant.config_entries import ConfigFlow
@@ -188,3 +190,70 @@ async def test_reloading_an_entry_reopens_the_connection(
     [second] = hass.data[DATA_MODBUS_CONNECTIONS].values()
 
     assert second.connection is not first.connection
+
+
+async def test_a_temporary_unit_closes_the_connection_on_exit(
+    hass: HomeAssistant,
+) -> None:
+    """A config flow's hold ends with the context, not with a config entry."""
+    with patch.object(ModbusConnection, "close") as close:
+        async with async_get_temporary_unit(
+            hass, ModbusTcpParams(host="1.2.3.4", port=502), 1
+        ):
+            [shared] = hass.data[DATA_MODBUS_CONNECTIONS].values()
+            assert shared.consumers == 1
+
+    assert close.called
+    assert not hass.data[DATA_MODBUS_CONNECTIONS]
+
+
+async def test_a_temporary_unit_releases_when_the_context_raises(
+    hass: HomeAssistant,
+) -> None:
+    """A flow step failing must not leak the connection it probed over."""
+    with patch.object(ModbusConnection, "close") as close, pytest.raises(ValueError):
+        async with async_get_temporary_unit(
+            hass, ModbusTcpParams(host="1.2.3.4", port=502), 1
+        ):
+            raise ValueError
+
+    assert close.called
+    assert not hass.data[DATA_MODBUS_CONNECTIONS]
+
+
+async def test_a_temporary_unit_shares_a_connection_an_entry_holds(
+    hass: HomeAssistant, consumer: ConsumerFactory
+) -> None:
+    """A flow probing a device an entry already talks to joins its connection.
+
+    The connection outlives the flow because the entry still holds it.
+    """
+    entry = consumer()
+    await hass.config_entries.async_setup(entry.entry_id)
+    params = ModbusTcpParams(host="1.2.3.4", port=502)
+    async_get_unit(hass, entry, params, 1)
+    [shared] = hass.data[DATA_MODBUS_CONNECTIONS].values()
+
+    async with async_get_temporary_unit(hass, params, 2):
+        assert shared.consumers == 2
+
+    assert shared.consumers == 1
+    assert hass.data[DATA_MODBUS_CONNECTIONS]
+
+
+async def test_a_temporary_unit_cannot_clash_with_held_link_settings(
+    hass: HomeAssistant, consumer: ConsumerFactory
+) -> None:
+    """A flow gets told about a link settings clash when entering the context."""
+    entry = consumer()
+    await hass.config_entries.async_setup(entry.entry_id)
+    async_get_unit(hass, entry, ModbusTcpParams(host="1.2.3.4", port=502), 1)
+
+    with pytest.raises(HomeAssistantError, match="different link settings"):
+        async with async_get_temporary_unit(
+            hass, ModbusTcpParams(host="1.2.3.4", port=502, framer="rtu"), 2
+        ):
+            pass
+
+    [shared] = hass.data[DATA_MODBUS_CONNECTIONS].values()
+    assert shared.consumers == 1