From 6f097eecc3b64e2c6feb26ec6c36e21a4e8cc71b Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Fri, 3 Feb 2023 21:22:31 +0100 Subject: [PATCH] Add WS commands thread/list_datasets, thread/get_dataset_tlv (#87333) --- .../components/thread/dataset_store.py | 24 +++++- .../components/thread/websocket_api.py | 56 +++++++++++++ tests/components/thread/test_dataset_store.py | 49 +++++++++-- tests/components/thread/test_websocket_api.py | 83 ++++++++++++++++++- 4 files changed, 203 insertions(+), 9 deletions(-) diff --git a/homeassistant/components/thread/dataset_store.py b/homeassistant/components/thread/dataset_store.py index a8fbd58d0d0..d19828d14d8 100644 --- a/homeassistant/components/thread/dataset_store.py +++ b/homeassistant/components/thread/dataset_store.py @@ -36,6 +36,21 @@ class DatasetEntry: """Return the dataset in dict format.""" return tlv_parser.parse_tlv(self.tlv) + @property + def extended_pan_id(self) -> str | None: + """Return extended PAN ID as a hex string.""" + return self.dataset.get(tlv_parser.MeshcopTLVType.EXTPANID) + + @property + def network_name(self) -> str | None: + """Return network name as a string.""" + return self.dataset.get(tlv_parser.MeshcopTLVType.NETWORKNAME) + + @property + def pan_id(self) -> str | None: + """Return PAN ID as a hex string.""" + return self.dataset.get(tlv_parser.MeshcopTLVType.PANID) + def to_json(self) -> dict[str, Any]: """Return a JSON serializable representation for storage.""" return { @@ -77,6 +92,11 @@ class DatasetStore: self.datasets[entry.id] = entry self.async_schedule_save() + @callback + def async_get(self, dataset_id: str) -> DatasetEntry | None: + """Get dataset by id.""" + return self.datasets.get(dataset_id) + async def async_load(self) -> None: """Load the datasets.""" data = await self._store.async_load() @@ -110,7 +130,7 @@ class DatasetStore: @singleton(DATA_STORE) -async def _async_get_store(hass: HomeAssistant) -> DatasetStore: +async def async_get_store(hass: HomeAssistant) -> DatasetStore: """Get the dataset store.""" store = DatasetStore(hass) await store.async_load() @@ -119,5 +139,5 @@ async def _async_get_store(hass: HomeAssistant) -> DatasetStore: async def async_add_dataset(hass: HomeAssistant, source: str, tlv: str) -> None: """Add a dataset.""" - store = await _async_get_store(hass) + store = await async_get_store(hass) store.async_add(source, tlv) diff --git a/homeassistant/components/thread/websocket_api.py b/homeassistant/components/thread/websocket_api.py index 3da8ec26420..82ac2a7e4c1 100644 --- a/homeassistant/components/thread/websocket_api.py +++ b/homeassistant/components/thread/websocket_api.py @@ -16,6 +16,8 @@ from . import dataset_store def async_setup(hass: HomeAssistant) -> None: """Set up the sensor websocket API.""" websocket_api.async_register_command(hass, ws_add_dataset) + websocket_api.async_register_command(hass, ws_get_dataset) + websocket_api.async_register_command(hass, ws_list_datasets) @websocket_api.require_admin @@ -43,3 +45,57 @@ async def ws_add_dataset( return connection.send_result(msg["id"]) + + +@websocket_api.require_admin +@websocket_api.websocket_command( + { + vol.Required("type"): "thread/get_dataset_tlv", + vol.Required("dataset_id"): str, + } +) +@websocket_api.async_response +async def ws_get_dataset( + hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] +) -> None: + """Get a thread dataset in TLV format.""" + dataset_id = msg["dataset_id"] + + store = await dataset_store.async_get_store(hass) + if not (dataset := store.async_get(dataset_id)): + connection.send_error( + msg["id"], websocket_api.const.ERR_NOT_FOUND, "unknown dataset" + ) + return + + connection.send_result(msg["id"], {"tlv": dataset.tlv}) + + +@websocket_api.require_admin +@websocket_api.websocket_command( + { + vol.Required("type"): "thread/list_datasets", + } +) +@websocket_api.async_response +async def ws_list_datasets( + hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] +) -> None: + """Get a list of thread datasets.""" + + store = await dataset_store.async_get_store(hass) + result = [] + for dataset in store.datasets.values(): + result.append( + { + "created": dataset.created, + "dataset_id": dataset.id, + "extended_pan_id": dataset.extended_pan_id, + "network_name": dataset.network_name, + "pan_id": dataset.pan_id, + "preferred": dataset.preferred, + "source": dataset.source, + } + ) + + connection.send_result(msg["id"], {"datasets": result}) diff --git a/tests/components/thread/test_dataset_store.py b/tests/components/thread/test_dataset_store.py index 02fd90124f8..483a3f68080 100644 --- a/tests/components/thread/test_dataset_store.py +++ b/tests/components/thread/test_dataset_store.py @@ -22,7 +22,7 @@ async def test_add_invalid_dataset(hass: HomeAssistant) -> None: with pytest.raises(TLVError, match="unknown type 222"): await dataset_store.async_add_dataset(hass, "source", "DEADBEEF") - store = await dataset_store._async_get_store(hass) + store = await dataset_store.async_get_store(hass) assert len(store.datasets) == 0 @@ -30,7 +30,7 @@ async def test_add_dataset_twice(hass: HomeAssistant) -> None: """Test adding dataset twice does nothing.""" await dataset_store.async_add_dataset(hass, "source", DATASET_1) - store = await dataset_store._async_get_store(hass) + store = await dataset_store.async_get_store(hass) assert len(store.datasets) == 1 created = list(store.datasets.values())[0].created @@ -43,7 +43,7 @@ async def test_add_dataset_reordered(hass: HomeAssistant) -> None: """Test adding dataset with keys in a different order does nothing.""" await dataset_store.async_add_dataset(hass, "source", DATASET_1) - store = await dataset_store._async_get_store(hass) + store = await dataset_store.async_get_store(hass) assert len(store.datasets) == 1 created = list(store.datasets.values())[0].created @@ -52,6 +52,45 @@ async def test_add_dataset_reordered(hass: HomeAssistant) -> None: assert list(store.datasets.values())[0].created == created +async def test_dataset_properties(hass: HomeAssistant) -> None: + """Test dataset entry properties.""" + datasets = [ + {"source": "Google", "tlv": DATASET_1}, + {"source": "Multipan", "tlv": DATASET_2}, + {"source": "🎅", "tlv": DATASET_3}, + ] + + for dataset in datasets: + await dataset_store.async_add_dataset(hass, dataset["source"], dataset["tlv"]) + + store = await dataset_store.async_get_store(hass) + for dataset in store.datasets.values(): + if dataset.source == "Google": + dataset_1 = dataset + if dataset.source == "Multipan": + dataset_2 = dataset + if dataset.source == "🎅": + dataset_3 = dataset + + dataset = store.async_get(dataset_1.id) + assert dataset == dataset_1 + assert dataset.extended_pan_id == "1111111122222222" + assert dataset.network_name == "OpenThreadDemo" + assert dataset.pan_id == "1234" + + dataset = store.async_get(dataset_2.id) + assert dataset == dataset_2 + assert dataset.extended_pan_id == "1111111122222222" + assert dataset.network_name == "HomeAssistant!" + assert dataset.pan_id == "1234" + + dataset = store.async_get(dataset_3.id) + assert dataset == dataset_3 + assert dataset.extended_pan_id == "1111111122222222" + assert dataset.network_name == "~🐣🐥🐤~" + assert dataset.pan_id == "1234" + + async def test_load_datasets(hass: HomeAssistant) -> None: """Make sure that we can load/save data correctly.""" @@ -70,7 +109,7 @@ async def test_load_datasets(hass: HomeAssistant) -> None: }, ] - store1 = await dataset_store._async_get_store(hass) + store1 = await dataset_store.async_get_store(hass) for dataset in datasets: store1.async_add(dataset["source"], dataset["tlv"]) assert len(store1.datasets) == 3 @@ -137,5 +176,5 @@ async def test_loading_datasets_from_storage(hass: HomeAssistant, hass_storage) }, } - store = await dataset_store._async_get_store(hass) + store = await dataset_store.async_get_store(hass) assert len(store.datasets) == 3 diff --git a/tests/components/thread/test_websocket_api.py b/tests/components/thread/test_websocket_api.py index 8955e876c4d..10c37258f93 100644 --- a/tests/components/thread/test_websocket_api.py +++ b/tests/components/thread/test_websocket_api.py @@ -5,7 +5,7 @@ from homeassistant.components.thread.const import DOMAIN from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component -from . import DATASET_1 +from . import DATASET_1, DATASET_2, DATASET_3 async def test_add_dataset(hass: HomeAssistant, hass_ws_client) -> None: @@ -22,7 +22,7 @@ async def test_add_dataset(hass: HomeAssistant, hass_ws_client) -> None: assert msg["success"] assert msg["result"] is None - store = await dataset_store._async_get_store(hass) + store = await dataset_store.async_get_store(hass) assert len(store.datasets) == 1 dataset = next(iter(store.datasets.values())) assert dataset.source == "test" @@ -42,3 +42,82 @@ async def test_add_invalid_dataset(hass: HomeAssistant, hass_ws_client) -> None: msg = await client.receive_json() assert not msg["success"] assert msg["error"] == {"code": "invalid_format", "message": "unknown type 222"} + + +async def test_list_get_dataset(hass: HomeAssistant, hass_ws_client) -> None: + """Test list and get datasets.""" + assert await async_setup_component(hass, DOMAIN, {}) + await hass.async_block_till_done() + + client = await hass_ws_client(hass) + + await client.send_json({"id": 1, "type": "thread/list_datasets"}) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == {"datasets": []} + + datasets = [ + {"source": "Google", "tlv": DATASET_1}, + {"source": "Multipan", "tlv": DATASET_2}, + {"source": "🎅", "tlv": DATASET_3}, + ] + for dataset in datasets: + await dataset_store.async_add_dataset(hass, dataset["source"], dataset["tlv"]) + + store = await dataset_store.async_get_store(hass) + for dataset in store.datasets.values(): + if dataset.source == "Google": + dataset_1 = dataset + if dataset.source == "Multipan": + dataset_2 = dataset + if dataset.source == "🎅": + dataset_3 = dataset + + await client.send_json({"id": 2, "type": "thread/list_datasets"}) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == { + "datasets": [ + { + "created": dataset_1.created.isoformat(), + "dataset_id": dataset_1.id, + "extended_pan_id": "1111111122222222", + "network_name": "OpenThreadDemo", + "pan_id": "1234", + "preferred": True, + "source": "Google", + }, + { + "created": dataset_2.created.isoformat(), + "dataset_id": dataset_2.id, + "extended_pan_id": "1111111122222222", + "network_name": "HomeAssistant!", + "pan_id": "1234", + "preferred": False, + "source": "Multipan", + }, + { + "created": dataset_3.created.isoformat(), + "dataset_id": dataset_3.id, + "extended_pan_id": "1111111122222222", + "network_name": "~🐣🐥🐤~", + "pan_id": "1234", + "preferred": False, + "source": "🎅", + }, + ] + } + + await client.send_json( + {"id": 3, "type": "thread/get_dataset_tlv", "dataset_id": dataset_2.id} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == {"tlv": dataset_2.tlv} + + await client.send_json( + {"id": 4, "type": "thread/get_dataset_tlv", "dataset_id": "blah"} + ) + msg = await client.receive_json() + assert not msg["success"] + assert msg["error"] == {"code": "not_found", "message": "unknown dataset"}