jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-04-23 08:47:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add DHCP discovery subscribe websocket API (#143106)

Browse Source 
* Add DHCP discovery subscribe websocket API

* fix circular import

* fixes

* fixes

* fixes

* reduce

* reduce

* reduce

* fix tests

* fix tests

* rework

* tests

* reduce number of lines changed

* reduce
This commit is contained in:
J. Nick Koston 2025-04-21 04:25:04 -10:00 committed by GitHub
parent 4b8447bc82
commit ba6ce28d3c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 381 additions and 158 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
homeassistant/components/dhcp/__init__.py
View File

@ -4,7 +4,6 @@ from __future__ import annotations
import asyncio
from collections.abc import Callable
from dataclasses import dataclass
from datetime import timedelta
from fnmatch import translate
from functools import lru_cache, partial
@ -66,13 +65,12 @@ from homeassistant.helpers.service_info.dhcp import DhcpServiceInfo as _DhcpServ
from homeassistant.helpers.typing import ConfigType
from homeassistant.loader import DHCPMatcher, async_get_dhcp
from .const import DOMAIN
from . import websocket_api
from .const import DOMAIN, HOSTNAME, IP_ADDRESS, MAC_ADDRESS
from .models import DATA_DHCP, DHCPAddressData, DHCPData, DhcpMatchers
CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)
HOSTNAME: Final = "hostname"
MAC_ADDRESS: Final = "macaddress"
IP_ADDRESS: Final = "ip"
REGISTERED_DEVICES: Final = "registered_devices"
SCAN_INTERVAL = timedelta(minutes=60)
@ -87,15 +85,6 @@ _DEPRECATED_DhcpServiceInfo = DeprecatedConstant(
)
@dataclass(slots=True)
class DhcpMatchers:
"""Prepared info from dhcp entries."""
registered_devices_domains: set[str]
no_oui_matchers: dict[str, list[DHCPMatcher]]
oui_matchers: dict[str, list[DHCPMatcher]]
def async_index_integration_matchers(
integration_matchers: list[DHCPMatcher],
) -> DhcpMatchers:
@ -133,36 +122,34 @@ def async_index_integration_matchers(
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up the dhcp component."""
watchers: list[WatcherBase] = []
address_data: dict[str, dict[str, str]] = {}
integration_matchers = async_index_integration_matchers(await async_get_dhcp(hass))
dhcp_data = DHCPData(integration_matchers=integration_matchers)
hass.data[DATA_DHCP] = dhcp_data
websocket_api.async_setup(hass)
watchers: list[WatcherBase] = []
# For the passive classes we need to start listening
# for state changes and connect the dispatchers before
# everything else starts up or we will miss events
device_watcher = DeviceTrackerWatcher(hass, address_data, integration_matchers)
device_watcher = DeviceTrackerWatcher(hass, dhcp_data)
device_watcher.async_start()
watchers.append(device_watcher)
device_tracker_registered_watcher = DeviceTrackerRegisteredWatcher(
hass, address_data, integration_matchers
)
device_tracker_registered_watcher = DeviceTrackerRegisteredWatcher(hass, dhcp_data)
device_tracker_registered_watcher.async_start()
watchers.append(device_tracker_registered_watcher)
async def _async_initialize(event: Event) -> None:
await aiodhcpwatcher.async_init()
network_watcher = NetworkWatcher(hass, address_data, integration_matchers)
network_watcher = NetworkWatcher(hass, dhcp_data)
network_watcher.async_start()
watchers.append(network_watcher)
dhcp_watcher = DHCPWatcher(hass, address_data, integration_matchers)
dhcp_watcher = DHCPWatcher(hass, dhcp_data)
await dhcp_watcher.async_start()
watchers.append(dhcp_watcher)
rediscovery_watcher = RediscoveryWatcher(
hass, address_data, integration_matchers
)
rediscovery_watcher = RediscoveryWatcher(hass, dhcp_data)
rediscovery_watcher.async_start()
watchers.append(rediscovery_watcher)
@ -180,18 +167,13 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
class WatcherBase:
"""Base class for dhcp and device tracker watching."""
def __init__(
self,
hass: HomeAssistant,
address_data: dict[str, dict[str, str]],
integration_matchers: DhcpMatchers,
) -> None:
def __init__(self, hass: HomeAssistant, dhcp_data: DHCPData) -> None:
"""Initialize class."""
super().__init__()
self.hass = hass
self._integration_matchers = integration_matchers
self._address_data = address_data
self._callbacks = dhcp_data.callbacks
self._integration_matchers = dhcp_data.integration_matchers
self._address_data = dhcp_data.address_data
self._unsub: Callable[[], None] | None = None
@callback
@ -230,18 +212,18 @@ class WatcherBase:
mac_address = formatted_mac.replace(":", "")
compressed_ip_address = made_ip_address.compressed
data = self._address_data.get(mac_address)
current_data = self._address_data.get(mac_address)
if (
not force
and data
and data[IP_ADDRESS] == compressed_ip_address
and data[HOSTNAME].startswith(hostname)
and current_data
and current_data[IP_ADDRESS] == compressed_ip_address
and current_data[HOSTNAME].startswith(hostname)
):
# If the address data is the same no need
# to process it
return
data = {IP_ADDRESS: compressed_ip_address, HOSTNAME: hostname}
data: DHCPAddressData = {IP_ADDRESS: compressed_ip_address, HOSTNAME: hostname}
self._address_data[mac_address] = data
lowercase_hostname = hostname.lower()
@ -287,9 +269,19 @@ class WatcherBase:
_LOGGER.debug("Matched %s against %s", data, matcher)
matched_domains.add(domain)
if not matched_domains:
return # avoid creating DiscoveryKey if there are no matches
if self._callbacks:
address_data = {mac_address: data}
for callback_ in self._callbacks:
callback_(address_data)
service_info: _DhcpServiceInfo | None = None
if not matched_domains:
return
service_info = _DhcpServiceInfo(
ip=ip_address,
hostname=lowercase_hostname,
macaddress=mac_address,
)
discovery_key = DiscoveryKey(
domain=DOMAIN,
key=mac_address,
@ -300,11 +292,7 @@ class WatcherBase:
self.hass,
domain,
{"source": config_entries.SOURCE_DHCP},
_DhcpServiceInfo(
ip=ip_address,
hostname=lowercase_hostname,
macaddress=mac_address,
),
service_info,
discovery_key=discovery_key,
)
@ -315,11 +303,10 @@ class NetworkWatcher(WatcherBase):
def __init__(
self,
hass: HomeAssistant,
address_data: dict[str, dict[str, str]],
integration_matchers: DhcpMatchers,
dhcp_data: DHCPData,
) -> None:
"""Initialize class."""
super().__init__(hass, address_data, integration_matchers)
super().__init__(hass, dhcp_data)
self._discover_hosts: DiscoverHosts | None = None
self._discover_task: asyncio.Task | None = None

5
homeassistant/components/dhcp/const.py
View File

@ -1,3 +1,8 @@
"""Constants for the dhcp integration."""
from typing import Final
DOMAIN = "dhcp"
HOSTNAME: Final = "hostname"
MAC_ADDRESS: Final = "macaddress"
IP_ADDRESS: Final = "ip"

37
homeassistant/components/dhcp/helpers.py Normal file
View File

@ -0,0 +1,37 @@
"""The dhcp integration."""
from __future__ import annotations
from collections.abc import Callable
from functools import partial
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
from .models import DATA_DHCP, DHCPAddressData
@callback
def async_register_dhcp_callback_internal(
hass: HomeAssistant,
callback_: Callable[[dict[str, DHCPAddressData]], None],
) -> CALLBACK_TYPE:
"""Register a dhcp callback.
For internal use only.
This is not intended for use by integrations.
"""
callbacks = hass.data[DATA_DHCP].callbacks
callbacks.add(callback_)
return partial(callbacks.remove, callback_)
@callback
def async_get_address_data_internal(
hass: HomeAssistant,
) -> dict[str, DHCPAddressData]:
"""Get the address data.
For internal use only.
This is not intended for use by integrations.
"""
return hass.data[DATA_DHCP].address_data

43
homeassistant/components/dhcp/models.py Normal file
View File

@ -0,0 +1,43 @@
"""The dhcp integration."""
from __future__ import annotations
from collections.abc import Callable
import dataclasses
from dataclasses import dataclass
from typing import TypedDict
from homeassistant.loader import DHCPMatcher
from homeassistant.util.hass_dict import HassKey
from .const import DOMAIN
@dataclass(slots=True)
class DhcpMatchers:
"""Prepared info from dhcp entries."""
registered_devices_domains: set[str]
no_oui_matchers: dict[str, list[DHCPMatcher]]
oui_matchers: dict[str, list[DHCPMatcher]]
class DHCPAddressData(TypedDict):
"""Typed dict for DHCP address data."""
hostname: str
ip: str
@dataclasses.dataclass(slots=True)
class DHCPData:
"""Data for the dhcp component."""
integration_matchers: DhcpMatchers
callbacks: set[Callable[[dict[str, DHCPAddressData]], None]] = dataclasses.field(
default_factory=set
)
address_data: dict[str, DHCPAddressData] = dataclasses.field(default_factory=dict)
DATA_DHCP: HassKey[DHCPData] = HassKey(DOMAIN)

63
homeassistant/components/dhcp/websocket_api.py Normal file
View File

@ -0,0 +1,63 @@
"""The dhcp integration websocket apis."""
from __future__ import annotations
from typing import Any
import voluptuous as vol
from homeassistant.components import websocket_api
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.json import json_bytes
from .const import HOSTNAME, IP_ADDRESS
from .helpers import (
async_get_address_data_internal,
async_register_dhcp_callback_internal,
)
from .models import DHCPAddressData
@callback
def async_setup(hass: HomeAssistant) -> None:
"""Set up the DHCP websocket API."""
websocket_api.async_register_command(hass, ws_subscribe_discovery)
@websocket_api.require_admin
@websocket_api.websocket_command(
{
vol.Required("type"): "dhcp/subscribe_discovery",
}
)
@websocket_api.async_response
async def ws_subscribe_discovery(
hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
) -> None:
"""Handle subscribe discovery websocket command."""
ws_msg_id: int = msg["id"]
def _async_send(address_data: dict[str, DHCPAddressData]) -> None:
connection.send_message(
json_bytes(
websocket_api.event_message(
ws_msg_id,
{
"add": [
{
"mac_address": dr.format_mac(mac_address).upper(),
"hostname": data[HOSTNAME],
"ip_address": data[IP_ADDRESS],
}
for mac_address, data in address_data.items()
]
},
)
)
)
unsub = async_register_dhcp_callback_internal(hass, _async_send)
connection.subscriptions[ws_msg_id] = unsub
connection.send_message(json_bytes(websocket_api.result_message(ws_msg_id)))
_async_send(async_get_address_data_internal(hass))

231
tests/components/dhcp/test_init.py
View File

@ -1,5 +1,7 @@
"""Test the DHCP discovery integration."""
from __future__ import annotations
from collections.abc import Awaitable, Callable
import datetime
import threading
@ -24,6 +26,7 @@ from homeassistant.components.device_tracker import (
SourceType,
)
from homeassistant.components.dhcp.const import DOMAIN
from homeassistant.components.dhcp.models import DHCPData
from homeassistant.const import (
EVENT_HOMEASSISTANT_STARTED,
EVENT_HOMEASSISTANT_STOP,
@ -147,12 +150,12 @@ async def _async_get_handle_dhcp_packet(
integration_matchers: dhcp.DhcpMatchers,
address_data: dict | None = None,
) -> Callable[[Any], Awaitable[None]]:
"""Make a handler for a dhcp packet."""
if address_data is None:
address_data = {}
dhcp_watcher = dhcp.DHCPWatcher(
hass,
address_data,
integration_matchers,
DHCPData(integration_matchers, set(), address_data),
)
with patch("aiodhcpwatcher.async_start"):
await dhcp_watcher.async_start()
@ -666,6 +669,45 @@ async def test_setup_fails_with_broken_libpcap(
)
def _make_device_tracker_watcher(
hass: HomeAssistant, matchers: list[dhcp.DHCPMatcher]
) -> dhcp.DeviceTrackerWatcher:
return dhcp.DeviceTrackerWatcher(
hass,
DHCPData(
dhcp.async_index_integration_matchers(matchers),
set(),
{},
),
)
def _make_device_tracker_registered_watcher(
hass: HomeAssistant, matchers: list[dhcp.DHCPMatcher]
) -> dhcp.DeviceTrackerRegisteredWatcher:
return dhcp.DeviceTrackerRegisteredWatcher(
hass,
DHCPData(
dhcp.async_index_integration_matchers(matchers),
set(),
{},
),
)
def _make_network_watcher(
hass: HomeAssistant, matchers: list[dhcp.DHCPMatcher]
) -> dhcp.NetworkWatcher:
return dhcp.NetworkWatcher(
hass,
DHCPData(
dhcp.async_index_integration_matchers(matchers),
set(),
{},
),
)
async def test_device_tracker_hostname_and_macaddress_exists_before_start(
hass: HomeAssistant,
) -> None:
@ -682,18 +724,15 @@ async def test_device_tracker_hostname_and_macaddress_exists_before_start(
)
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -716,18 +755,15 @@ async def test_device_tracker_hostname_and_macaddress_exists_before_start(
async def test_device_tracker_registered(hass: HomeAssistant) -> None:
"""Test matching based on hostname and macaddress when registered."""
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher(
device_tracker_watcher = _make_device_tracker_registered_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -756,18 +792,15 @@ async def test_device_tracker_registered(hass: HomeAssistant) -> None:
async def test_device_tracker_registered_hostname_none(hass: HomeAssistant) -> None:
"""Test handle None hostname."""
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerRegisteredWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -789,18 +822,15 @@ async def test_device_tracker_hostname_and_macaddress_after_start(
"""Test matching based on hostname and macaddress after start."""
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -837,18 +867,15 @@ async def test_device_tracker_hostname_and_macaddress_after_start_not_home(
"""Test matching based on hostname and macaddress after start but not home."""
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -875,9 +902,8 @@ async def test_device_tracker_hostname_and_macaddress_after_start_not_router(
"""Test matching based on hostname and macaddress after start but not router."""
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
)
device_tracker_watcher.async_start()
@ -905,9 +931,8 @@ async def test_device_tracker_hostname_and_macaddress_after_start_hostname_missi
"""Test matching based on hostname and macaddress after start but missing hostname."""
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
)
device_tracker_watcher.async_start()
@ -934,9 +959,8 @@ async def test_device_tracker_invalid_ip_address(
"""Test an invalid ip address."""
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
[{"domain": "mock-domain", "hostname": "connect", "macaddress": "B8B7F1*"}],
)
device_tracker_watcher.async_start()
@ -974,18 +998,15 @@ async def test_device_tracker_ignore_self_assigned_ips_before_start(
)
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
device_tracker_watcher = dhcp.DeviceTrackerWatcher(
device_tracker_watcher = _make_device_tracker_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -1010,18 +1031,15 @@ async def test_aiodiscover_finds_new_hosts(hass: HomeAssistant) -> None:
],
),
):
device_tracker_watcher = dhcp.NetworkWatcher(
device_tracker_watcher = _make_network_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -1073,18 +1091,15 @@ async def test_aiodiscover_does_not_call_again_on_shorter_hostname(
],
),
):
device_tracker_watcher = dhcp.NetworkWatcher(
device_tracker_watcher = _make_network_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "irobot-*",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "irobot-*",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -1123,19 +1138,17 @@ async def test_aiodiscover_finds_new_hosts_after_interval(hass: HomeAssistant) -
return_value=[],
),
):
device_tracker_watcher = dhcp.NetworkWatcher(
device_tracker_watcher = _make_network_watcher(
hass,
{},
dhcp.async_index_integration_matchers(
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
]
),
[
{
"domain": "mock-domain",
"hostname": "connect",
"macaddress": "B8B7F1*",
}
],
)
device_tracker_watcher.async_start()
await hass.async_block_till_done()
@ -1235,7 +1248,7 @@ async def test_dhcp_rediscover(
hass, integration_matchers, address_data
)
rediscovery_watcher = dhcp.RediscoveryWatcher(
hass, address_data, integration_matchers
hass, DHCPData(integration_matchers, set(), address_data)
)
rediscovery_watcher.async_start()
with patch.object(hass.config_entries.flow, "async_init") as mock_init:
@ -1329,7 +1342,7 @@ async def test_dhcp_rediscover_no_match(
hass, integration_matchers, address_data
)
rediscovery_watcher = dhcp.RediscoveryWatcher(
hass, address_data, integration_matchers
hass, DHCPData(integration_matchers, set(), address_data)
)
rediscovery_watcher.async_start()
with patch.object(hass.config_entries.flow, "async_init") as mock_init:

75
tests/components/dhcp/test_websocket_api.py Normal file
View File

@ -0,0 +1,75 @@
"""The tests for the dhcp WebSocket API."""
import asyncio
from collections.abc import Callable
from unittest.mock import patch
import aiodhcpwatcher
from homeassistant.components.dhcp import DOMAIN
from homeassistant.core import EVENT_HOMEASSISTANT_STARTED, HomeAssistant
from homeassistant.setup import async_setup_component
from tests.typing import WebSocketGenerator
async def test_subscribe_discovery(
hass: HomeAssistant,
hass_ws_client: WebSocketGenerator,
) -> None:
"""Test dhcp subscribe_discovery."""
saved_callback: Callable[[aiodhcpwatcher.DHCPRequest], None] | None = None
async def mock_start(
callback: Callable[[aiodhcpwatcher.DHCPRequest], None],
) -> None:
"""Mock start."""
nonlocal saved_callback
saved_callback = callback
with (
patch("homeassistant.components.dhcp.aiodhcpwatcher.async_start", mock_start),
patch("homeassistant.components.dhcp.DiscoverHosts"),
):
await async_setup_component(hass, DOMAIN, {})
await hass.async_block_till_done()
hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
await hass.async_block_till_done()
saved_callback(aiodhcpwatcher.DHCPRequest("4.3.2.2", "happy", "44:44:33:11:23:12"))
client = await hass_ws_client()
await client.send_json(
{
"id": 1,
"type": "dhcp/subscribe_discovery",
}
)
async with asyncio.timeout(1):
response = await client.receive_json()
assert response["success"]
async with asyncio.timeout(1):
response = await client.receive_json()
assert response["event"] == {
"add": [
{
"hostname": "happy",
"ip_address": "4.3.2.2",
"mac_address": "44:44:33:11:23:12",
}
]
}
saved_callback(aiodhcpwatcher.DHCPRequest("4.3.2.1", "sad", "44:44:33:11:23:13"))
async with asyncio.timeout(1):
response = await client.receive_json()
assert response["event"] == {
"add": [
{
"hostname": "sad",
"ip_address": "4.3.2.1",
"mac_address": "44:44:33:11:23:13",
}
]
}