Expose the Modbus connections over a websocket API (#179937)

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

homeassistant/components/modbus/__init__.py

@@ -4,6 +4,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.frame import ReportBehavior, report_usage
 from homeassistant.helpers.typing import ConfigType
 
+from . import websocket_api
 from .connection import async_get_temporary_unit, async_get_unit
 from .const import DATA_MODBUS_HUBS, DOMAIN
 from .modbus import ModbusHub, async_modbus_setup
@@ -45,6 +46,7 @@ def get_hub(hass: HomeAssistant, name: str) -> ModbusHub:
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up Modbus component."""
     async_setup_services(hass)
+    websocket_api.async_setup(hass)
 
     if DOMAIN not in config:
         return True

homeassistant/components/modbus/connection.py

@@ -2,7 +2,7 @@
 
 from collections.abc import AsyncIterator, Callable, Coroutine
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import logging
 from typing import Any
 
@@ -36,19 +36,44 @@ DATA_MODBUS_CONNECTIONS: HassKey[dict[ModbusEndpoint, _SharedConnection]] = Hass
 
 @dataclass
 class _SharedConnection:
-    """A connection and how many units are held on it."""
+    """A connection and the units held on it."""
 
     params: ModbusParams
     connection: ModbusConnection
-    consumers: int = 0
+    units: dict[str, set[int]] = field(default_factory=dict)
+    """The unit ids each config entry holds, keyed by entry id."""
+    transient: int = 0
+    """Holds with no config entry behind them, taken by a config flow."""
+
+    @property
+    def consumers(self) -> int:
+        """How many holds are on this connection."""
+        return sum(len(held) for held in self.units.values()) + self.transient
+
+
+@dataclass(frozen=True, kw_only=True)
+class ModbusConnectionInfo:
+    """A connection the integration is keeping open, and who is using it."""
+
+    endpoint: ModbusEndpoint
+    connected: bool
+    units: dict[str, list[int]]
+    """The unit ids each config entry holds, keyed by entry id."""
 
 
 @callback
 def _async_acquire(
-    hass: HomeAssistant, params: ModbusParams
+    hass: HomeAssistant,
+    params: ModbusParams,
+    entry_id: str | None,
+    unit_id: int,
 ) -> tuple[ModbusConnection, Callable[[], Coroutine[Any, Any, None]]]:
     """Take a hold on the connection these credentials describe.
 
+    A hold with no ``entry_id`` behind it is a config flow's, which keeps the
+    connection up without belonging to anything that could be shown as using
+    it.
+
     Raises `HomeAssistantError` if the device is already in use over different
     link settings, which cannot both be honoured on one connection.
     """
@@ -64,11 +89,19 @@ def _async_acquire(
             f"settings: {shared.params} against {params}"
         )
 
-    shared.consumers += 1
+    if entry_id is None:
+        shared.transient += 1
+    else:
+        shared.units.setdefault(entry_id, set()).add(unit_id)
 
     async def release() -> None:
         """Give up this hold, closing behind the last one."""
-        shared.consumers -= 1
+        if entry_id is None:
+            shared.transient -= 1
+        elif (held := shared.units.get(entry_id)) is not None:
+            held.discard(unit_id)
+            if not held:
+                del shared.units[entry_id]
         if shared.consumers or connections.get(endpoint) is not shared:
             return
         del connections[endpoint]
@@ -94,7 +127,7 @@ def async_get_unit(
     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)
+    connection, release = _async_acquire(hass, params, entry.entry_id, unit_id)
     entry.async_on_unload(release)
     return connection.for_unit(unit_id)
 
@@ -114,8 +147,25 @@ async def async_get_temporary_unit(
     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)
+    connection, release = _async_acquire(hass, params, None, unit_id)
     try:
         yield connection.for_unit(unit_id)
     finally:
         await release()
+
+
+@callback
+def async_get_connection_info(hass: HomeAssistant) -> list[ModbusConnectionInfo]:
+    """Return the connections the integration is keeping open.
+
+    One entry per physical device, naming the config entries holding units on
+    it. A device several integrations share appears once, with all of them.
+    """
+    return [
+        ModbusConnectionInfo(
+            endpoint=endpoint,
+            connected=shared.connection.connected,
+            units={entry_id: sorted(held) for entry_id, held in shared.units.items()},
+        )
+        for endpoint, shared in hass.data.get(DATA_MODBUS_CONNECTIONS, {}).items()
+    ]

homeassistant/components/modbus/websocket_api.py

@@ -0,0 +1,42 @@
+"""Websocket API exposing the Modbus connections the integration keeps open."""
+
+from typing import Any, Final
+
+import voluptuous as vol
+
+from homeassistant.components import websocket_api
+from homeassistant.core import HomeAssistant, callback
+
+from .connection import async_get_connection_info
+
+TYPE_LIST_CONNECTIONS: Final = "modbus/connections/list"
+
+
+@callback
+def async_setup(hass: HomeAssistant) -> None:
+    """Register the Modbus websocket commands."""
+    websocket_api.async_register_command(hass, websocket_list_connections)
+
+
+@websocket_api.require_admin
+@websocket_api.websocket_command({vol.Required("type"): TYPE_LIST_CONNECTIONS})
+@callback
+def websocket_list_connections(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
+    """List the connections, and which config entries hold units on each."""
+    connection.send_result(
+        msg["id"],
+        {
+            "connections": [
+                {
+                    "endpoint": list(info.endpoint),
+                    "connected": info.connected,
+                    "units": info.units,
+                }
+                for info in async_get_connection_info(hass)
+            ]
+        },
+    )

tests/components/modbus/test_connection.py

@@ -257,3 +257,56 @@ async def test_a_temporary_unit_cannot_clash_with_held_link_settings(
 
     [shared] = hass.data[DATA_MODBUS_CONNECTIONS].values()
     assert shared.consumers == 1
+
+
+async def test_one_entry_holding_the_same_unit_twice(
+    hass: HomeAssistant, consumer: ConsumerFactory
+) -> None:
+    """Two holds on one unit are one unit, and release together.
+
+    The registry records which units an entry holds, not how many times it
+    asked, so asking twice adds nothing to release twice.
+    """
+    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)
+    async_get_unit(hass, entry, params, 1)
+    [shared] = hass.data[DATA_MODBUS_CONNECTIONS].values()
+    assert shared.units == {entry.entry_id: {1}}
+
+    with patch.object(shared.connection, "close") as close:
+        await hass.config_entries.async_unload(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert close.call_count == 1
+    assert not hass.data[DATA_MODBUS_CONNECTIONS]
+
+
+async def test_two_entries_holding_the_same_unit(
+    hass: HomeAssistant, consumer: ConsumerFactory
+) -> None:
+    """The link stays while anybody still holds that unit, however they asked."""
+    one = consumer()
+    await hass.config_entries.async_setup(one.entry_id)
+    two = consumer()
+    await hass.config_entries.async_setup(two.entry_id)
+
+    params = ModbusTcpParams(host="1.2.3.4", port=502)
+    async_get_unit(hass, one, params, 1)
+    async_get_unit(hass, one, params, 1)  # the same unit, asked for twice
+    async_get_unit(hass, two, params, 1)
+    [shared] = hass.data[DATA_MODBUS_CONNECTIONS].values()
+
+    with patch.object(shared.connection, "close") as close:
+        await hass.config_entries.async_unload(one.entry_id)
+        await hass.async_block_till_done()
+
+        assert not close.called  # the other entry is still on that unit
+        assert hass.data[DATA_MODBUS_CONNECTIONS]
+
+        await hass.config_entries.async_unload(two.entry_id)
+        await hass.async_block_till_done()
+
+    assert close.called

tests/components/modbus/test_websocket_api.py

@@ -0,0 +1,187 @@
+"""Test the Modbus websocket API."""
+
+from collections.abc import Callable, Generator
+from unittest.mock import AsyncMock, patch
+
+from modbus_connection import ModbusTcpParams
+from modbus_connection.tmodbus import ModbusConnection
+import pytest
+
+from homeassistant.components.modbus import async_get_unit
+from homeassistant.config_entries import ConfigFlow
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from tests.common import (
+    MockConfigEntry,
+    MockModule,
+    mock_config_flow,
+    mock_integration,
+    mock_platform,
+)
+from tests.typing import WebSocketGenerator
+
+type ConsumerFactory = Callable[[], MockConfigEntry]
+
+
+class MockFlow(ConfigFlow):
+    """A config flow for the integration standing in for a consumer."""
+
+
+@pytest.fixture(name="consumer")
+def consumer_fixture(hass: HomeAssistant) -> Generator[ConsumerFactory]:
+    """Return a factory for config entries that can be set up and unloaded."""
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=AsyncMock(return_value=True),
+            async_unload_entry=AsyncMock(return_value=True),
+        ),
+    )
+    mock_platform(hass, "test.config_flow")
+
+    def _consumer() -> MockConfigEntry:
+        entry = MockConfigEntry(domain="test")
+        entry.add_to_hass(hass)
+        return entry
+
+    with mock_config_flow("test", MockFlow):
+        yield _consumer
+
+
+async def test_list_connections(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    consumer: ConsumerFactory,
+) -> None:
+    """Two entries on one device are one connection naming both."""
+    assert await async_setup_component(hass, "modbus", {})
+
+    first = consumer()
+    await hass.config_entries.async_setup(first.entry_id)
+    second = consumer()
+    await hass.config_entries.async_setup(second.entry_id)
+
+    params = ModbusTcpParams(host="device.local", port=502)
+    async_get_unit(hass, first, params, 1)
+    async_get_unit(hass, second, params, 2)
+
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id({"type": "modbus/connections/list"})
+    result = (await client.receive_json())["result"]
+
+    assert result == {
+        "connections": [
+            {
+                "endpoint": ["tcp", "device.local", 502],
+                "connected": False,
+                "units": {first.entry_id: [1], second.entry_id: [2]},
+            }
+        ]
+    }
+
+
+async def test_a_connection_that_is_up_reports_itself_connected(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    consumer: ConsumerFactory,
+) -> None:
+    """The reported state follows the connection, rather than being fixed."""
+    assert await async_setup_component(hass, "modbus", {})
+
+    entry = consumer()
+    await hass.config_entries.async_setup(entry.entry_id)
+    async_get_unit(hass, entry, ModbusTcpParams(host="device.local", port=502), 1)
+
+    client = await hass_ws_client(hass)
+    with patch.object(ModbusConnection, "connected", True):
+        await client.send_json_auto_id({"type": "modbus/connections/list"})
+        result = (await client.receive_json())["result"]
+
+    assert result == {
+        "connections": [
+            {
+                "endpoint": ["tcp", "device.local", 502],
+                "connected": True,
+                "units": {entry.entry_id: [1]},
+            }
+        ]
+    }
+
+
+async def test_listing_the_connections_requires_admin(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    hass_read_only_access_token: str,
+) -> None:
+    """The endpoint names devices and config entries, so admins only."""
+    assert await async_setup_component(hass, "modbus", {})
+
+    client = await hass_ws_client(hass, hass_read_only_access_token)
+    await client.send_json_auto_id({"type": "modbus/connections/list"})
+    response = await client.receive_json()
+
+    assert not response["success"]
+    assert response["error"]["code"] == "unauthorized"
+
+
+async def test_unloading_an_entry_drops_it_from_the_listing(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    consumer: ConsumerFactory,
+) -> None:
+    """The connection stays while somebody else holds a unit on it."""
+    assert await async_setup_component(hass, "modbus", {})
+
+    first = consumer()
+    await hass.config_entries.async_setup(first.entry_id)
+    second = consumer()
+    await hass.config_entries.async_setup(second.entry_id)
+
+    params = ModbusTcpParams(host="device.local", port=502)
+    async_get_unit(hass, first, params, 1)
+    async_get_unit(hass, second, params, 2)
+
+    await hass.config_entries.async_unload(first.entry_id)
+
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id({"type": "modbus/connections/list"})
+    result = (await client.receive_json())["result"]
+
+    assert len(result["connections"]) == 1
+    assert result["connections"][0]["units"] == {second.entry_id: [2]}
+
+
+async def test_no_connections_when_nobody_asked(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+) -> None:
+    """Nothing is opened until an integration asks for a unit."""
+    assert await async_setup_component(hass, "modbus", {})
+
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id({"type": "modbus/connections/list"})
+
+    assert (await client.receive_json())["result"] == {"connections": []}
+
+
+async def test_one_entry_holding_two_units(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    consumer: ConsumerFactory,
+) -> None:
+    """An entry with two devices on one link reports both units."""
+    assert await async_setup_component(hass, "modbus", {})
+
+    entry = consumer()
+    await hass.config_entries.async_setup(entry.entry_id)
+
+    params = ModbusTcpParams(host="device.local", port=502)
+    async_get_unit(hass, entry, params, 1)
+    async_get_unit(hass, entry, params, 2)
+
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id({"type": "modbus/connections/list"})
+    result = (await client.receive_json())["result"]
+
+    assert result["connections"][0]["units"] == {entry.entry_id: [1, 2]}