Allow registering a callback to ssdp that matches any key value (#51382)

This commit is contained in:
J. Nick Koston 2021-06-03 08:26:37 -10:00 committed by GitHub
parent 53ae340900
commit c1111afef8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 52 additions and 18 deletions

View File

@ -13,7 +13,11 @@ from async_upnp_client.utils import CaseInsensitiveDict
from homeassistant import config_entries
from homeassistant.components import network
from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EVENT_HOMEASSISTANT_STOP
from homeassistant.const import (
EVENT_HOMEASSISTANT_STARTED,
EVENT_HOMEASSISTANT_STOP,
MATCH_ALL,
)
from homeassistant.core import CoreState, HomeAssistant, callback as core_callback
from homeassistant.helpers.event import async_track_time_interval
from homeassistant.helpers.typing import ConfigType
@ -128,6 +132,19 @@ def _async_process_callbacks(
_LOGGER.exception("Failed to callback info: %s", discovery_info)
@core_callback
def _async_headers_match(
headers: Mapping[str, str], match_dict: dict[str, str]
) -> bool:
for header, val in match_dict.items():
if val == MATCH_ALL:
if header not in headers:
return False
elif headers.get(header) != val:
return False
return True
class Scanner:
"""Class to manage SSDP scanning."""
@ -157,7 +174,10 @@ class Scanner:
# before the callback was registered are fired
if self.hass.state != CoreState.running:
for headers in self.cache.values():
self._async_callback_if_match(callback, headers, match_dict)
if _async_headers_match(headers, match_dict):
_async_process_callbacks(
[callback], self._async_headers_to_discovery_info(headers)
)
callback_entry = (callback, match_dict)
self._callbacks.append(callback_entry)
@ -168,20 +188,6 @@ class Scanner:
return _async_remove_callback
@core_callback
def _async_callback_if_match(
self,
callback: Callable[[dict], None],
headers: Mapping[str, str],
match_dict: dict[str, str],
) -> None:
"""Fire a callback if info matches the match dict."""
if not all(headers.get(k) == v for (k, v) in match_dict.items()):
return
_async_process_callbacks(
[callback], self._async_headers_to_discovery_info(headers)
)
@core_callback
def async_stop(self, *_: Any) -> None:
"""Stop the scanner."""
@ -250,7 +256,7 @@ class Scanner:
return [
callback
for callback, match_dict in self._callbacks
if all(headers.get(k) == v for (k, v) in match_dict.items())
if _async_headers_match(headers, match_dict)
]
@core_callback

View File

@ -11,7 +11,11 @@ import pytest
from homeassistant import config_entries
from homeassistant.components import ssdp
from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EVENT_HOMEASSISTANT_STOP
from homeassistant.const import (
EVENT_HOMEASSISTANT_STARTED,
EVENT_HOMEASSISTANT_STOP,
MATCH_ALL,
)
from homeassistant.core import CoreState, callback
from homeassistant.setup import async_setup_component
import homeassistant.util.dt as dt_util
@ -356,9 +360,12 @@ async def test_scan_with_registered_callback(hass, aioclient_mock, caplog):
"location": "http://1.1.1.1",
"usn": "uuid:TIVRTLSR7ANF-D6E-1557809135086-RETAIL::urn:mdx-netflix-com:service:target:3",
"server": "mock-server",
"x-rincon-bootseq": "55",
"ext": "",
}
not_matching_intergration_callbacks = []
intergration_match_all_callbacks = []
intergration_match_all_not_present_callbacks = []
intergration_callbacks = []
intergration_callbacks_from_cache = []
match_any_callbacks = []
@ -371,6 +378,14 @@ async def test_scan_with_registered_callback(hass, aioclient_mock, caplog):
def _async_intergration_callbacks(info):
intergration_callbacks.append(info)
@callback
def _async_intergration_match_all_callbacks(info):
intergration_match_all_callbacks.append(info)
@callback
def _async_intergration_match_all_not_present_callbacks(info):
intergration_match_all_not_present_callbacks.append(info)
@callback
def _async_intergration_callbacks_from_cache(info):
intergration_callbacks_from_cache.append(info)
@ -410,6 +425,16 @@ async def test_scan_with_registered_callback(hass, aioclient_mock, caplog):
_async_intergration_callbacks,
{"st": "mock-st"},
)
ssdp.async_register_callback(
hass,
_async_intergration_match_all_callbacks,
{"x-rincon-bootseq": MATCH_ALL},
)
ssdp.async_register_callback(
hass,
_async_intergration_match_all_not_present_callbacks,
{"x-not-there": MATCH_ALL},
)
ssdp.async_register_callback(
hass,
_async_not_matching_intergration_callbacks,
@ -436,6 +461,8 @@ async def test_scan_with_registered_callback(hass, aioclient_mock, caplog):
assert len(intergration_callbacks) == 3
assert len(intergration_callbacks_from_cache) == 3
assert len(intergration_match_all_callbacks) == 3
assert len(intergration_match_all_not_present_callbacks) == 0
assert len(match_any_callbacks) == 3
assert len(not_matching_intergration_callbacks) == 0
assert intergration_callbacks[0] == {
@ -446,6 +473,7 @@ async def test_scan_with_registered_callback(hass, aioclient_mock, caplog):
ssdp.ATTR_SSDP_ST: "mock-st",
ssdp.ATTR_SSDP_USN: "uuid:TIVRTLSR7ANF-D6E-1557809135086-RETAIL::urn:mdx-netflix-com:service:target:3",
ssdp.ATTR_UPNP_UDN: "uuid:TIVRTLSR7ANF-D6E-1557809135086-RETAIL",
"x-rincon-bootseq": "55",
}
assert "Failed to callback info" in caplog.text