diff --git a/homeassistant/components/thread/__init__.py b/homeassistant/components/thread/__init__.py index dd2527763ad..679127e5202 100644 --- a/homeassistant/components/thread/__init__.py +++ b/homeassistant/components/thread/__init__.py @@ -11,7 +11,9 @@ from .dataset_store import ( DatasetEntry, async_add_dataset, async_get_dataset, + async_get_preferred_border_agent_id, async_get_preferred_dataset, + async_set_preferred_border_agent_id, ) from .websocket_api import async_setup as async_setup_ws_api @@ -19,8 +21,10 @@ __all__ = [ "DOMAIN", "DatasetEntry", "async_add_dataset", + "async_get_preferred_border_agent_id", "async_get_dataset", "async_get_preferred_dataset", + "async_set_preferred_border_agent_id", ] CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN) diff --git a/homeassistant/components/thread/dataset_store.py b/homeassistant/components/thread/dataset_store.py index 55623f7e3a4..96a9cf8e59e 100644 --- a/homeassistant/components/thread/dataset_store.py +++ b/homeassistant/components/thread/dataset_store.py @@ -19,7 +19,7 @@ from homeassistant.util import dt as dt_util, ulid as ulid_util DATA_STORE = "thread.datasets" STORAGE_KEY = "thread.datasets" STORAGE_VERSION_MAJOR = 1 -STORAGE_VERSION_MINOR = 2 +STORAGE_VERSION_MINOR = 3 SAVE_DELAY = 10 _LOGGER = logging.getLogger(__name__) @@ -86,7 +86,9 @@ class DatasetStoreStore(Store): ) -> dict[str, Any]: """Migrate to the new version.""" if old_major_version == 1: + data = old_data if old_minor_version < 2: + # Deduplicate datasets datasets: dict[str, DatasetEntry] = {} preferred_dataset = old_data["preferred_dataset"] @@ -156,6 +158,9 @@ class DatasetStoreStore(Store): "preferred_dataset": preferred_dataset, "datasets": [dataset.to_json() for dataset in datasets.values()], } + if old_minor_version < 3: + # Add border agent ID + data.setdefault("preferred_border_agent_id", None) return data @@ -167,6 +172,7 @@ class DatasetStore: """Initialize the dataset store.""" self.hass = hass self.datasets: dict[str, DatasetEntry] = {} + self._preferred_border_agent_id: str | None = None self._preferred_dataset: str | None = None self._store: Store[dict[str, Any]] = DatasetStoreStore( hass, @@ -259,6 +265,17 @@ class DatasetStore: """Get dataset by id.""" return self.datasets.get(dataset_id) + @callback + def async_get_preferred_border_agent_id(self) -> str | None: + """Get preferred border agent id.""" + return self._preferred_border_agent_id + + @callback + def async_set_preferred_border_agent_id(self, border_agent_id: str) -> None: + """Set preferred border agent id.""" + self._preferred_border_agent_id = border_agent_id + self.async_schedule_save() + @property @callback def preferred_dataset(self) -> str | None: @@ -279,6 +296,7 @@ class DatasetStore: data = await self._store.async_load() datasets: dict[str, DatasetEntry] = {} + preferred_border_agent_id: str | None = None preferred_dataset: str | None = None if data is not None: @@ -290,9 +308,11 @@ class DatasetStore: source=dataset["source"], tlv=dataset["tlv"], ) + preferred_border_agent_id = data["preferred_border_agent_id"] preferred_dataset = data["preferred_dataset"] self.datasets = datasets + self._preferred_border_agent_id = preferred_border_agent_id self._preferred_dataset = preferred_dataset @callback @@ -305,6 +325,7 @@ class DatasetStore: """Return data of datasets to store in a file.""" data: dict[str, Any] = {} data["datasets"] = [dataset.to_json() for dataset in self.datasets.values()] + data["preferred_border_agent_id"] = self._preferred_border_agent_id data["preferred_dataset"] = self._preferred_dataset return data @@ -331,6 +352,20 @@ async def async_get_dataset(hass: HomeAssistant, dataset_id: str) -> str | None: return entry.tlv +async def async_get_preferred_border_agent_id(hass: HomeAssistant) -> str | None: + """Get the preferred border agent ID.""" + store = await async_get_store(hass) + return store.async_get_preferred_border_agent_id() + + +async def async_set_preferred_border_agent_id( + hass: HomeAssistant, border_agent_id: str +) -> None: + """Get the preferred border agent ID.""" + store = await async_get_store(hass) + store.async_set_preferred_border_agent_id(border_agent_id) + + async def async_get_preferred_dataset(hass: HomeAssistant) -> str | None: """Get the preferred dataset.""" store = await async_get_store(hass) diff --git a/homeassistant/components/thread/websocket_api.py b/homeassistant/components/thread/websocket_api.py index 60941426b7e..853d8c3c893 100644 --- a/homeassistant/components/thread/websocket_api.py +++ b/homeassistant/components/thread/websocket_api.py @@ -20,6 +20,8 @@ def async_setup(hass: HomeAssistant) -> None: websocket_api.async_register_command(hass, ws_discover_routers) websocket_api.async_register_command(hass, ws_get_dataset) websocket_api.async_register_command(hass, ws_list_datasets) + websocket_api.async_register_command(hass, ws_get_preferred_border_agent_id) + websocket_api.async_register_command(hass, ws_set_preferred_border_agent_id) websocket_api.async_register_command(hass, ws_set_preferred_dataset) @@ -50,6 +52,38 @@ async def ws_add_dataset( connection.send_result(msg["id"]) +@websocket_api.require_admin +@websocket_api.websocket_command( + { + vol.Required("type"): "thread/get_preferred_border_agent_id", + } +) +@websocket_api.async_response +async def ws_get_preferred_border_agent_id( + hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] +) -> None: + """Get the preferred border agent ID.""" + border_agent_id = await dataset_store.async_get_preferred_border_agent_id(hass) + connection.send_result(msg["id"], {"border_agent_id": border_agent_id}) + + +@websocket_api.require_admin +@websocket_api.websocket_command( + { + vol.Required("type"): "thread/set_preferred_border_agent_id", + vol.Required("border_agent_id"): str, + } +) +@websocket_api.async_response +async def ws_set_preferred_border_agent_id( + hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] +) -> None: + """Set the preferred border agent ID.""" + border_agent_id = msg["border_agent_id"] + await dataset_store.async_set_preferred_border_agent_id(hass, border_agent_id) + connection.send_result(msg["id"]) + + @websocket_api.require_admin @websocket_api.websocket_command( { diff --git a/tests/components/thread/test_dataset_store.py b/tests/components/thread/test_dataset_store.py index 1ed754dbdcd..1171c597e99 100644 --- a/tests/components/thread/test_dataset_store.py +++ b/tests/components/thread/test_dataset_store.py @@ -319,12 +319,17 @@ async def test_loading_datasets_from_storage( "tlv": DATASET_3, }, ], + "preferred_border_agent_id": "230C6A1AC57F6F4BE262ACF32E5EF52C", "preferred_dataset": "id1", }, } store = await dataset_store.async_get_store(hass) assert len(store.datasets) == 3 + assert ( + store.async_get_preferred_border_agent_id() + == "230C6A1AC57F6F4BE262ACF32E5EF52C" + ) assert store.preferred_dataset == "id1" @@ -512,3 +517,34 @@ async def test_migrate_drop_duplicate_datasets_preferred( f"Dropped duplicated Thread dataset '{DATASET_1_LARGER_TIMESTAMP}' " f"(duplicate of preferred dataset '{DATASET_1}')" ) in caplog.text + + +async def test_migrate_set_default_border_agent_id( + hass: HomeAssistant, hass_storage: dict[str, Any], caplog +) -> None: + """Test migrating the dataset store adds default border agent.""" + hass_storage[dataset_store.STORAGE_KEY] = { + "version": 1, + "minor_version": 2, + "data": { + "datasets": [ + { + "created": "2023-02-02T09:41:13.746514+00:00", + "id": "id1", + "source": "source_1", + "tlv": DATASET_1, + }, + ], + "preferred_dataset": "id1", + }, + } + + store = await dataset_store.async_get_store(hass) + assert store.async_get_preferred_border_agent_id() is None + + +async def test_preferred_border_agent_id(hass: HomeAssistant) -> None: + """Test get and set the preferred border agent ID.""" + assert await dataset_store.async_get_preferred_border_agent_id(hass) is None + await dataset_store.async_set_preferred_border_agent_id(hass, "blah") + assert await dataset_store.async_get_preferred_border_agent_id(hass) == "blah" diff --git a/tests/components/thread/test_websocket_api.py b/tests/components/thread/test_websocket_api.py index 0db16318db1..82450474e92 100644 --- a/tests/components/thread/test_websocket_api.py +++ b/tests/components/thread/test_websocket_api.py @@ -200,6 +200,35 @@ async def test_list_get_dataset( assert msg["error"] == {"code": "not_found", "message": "unknown dataset"} +async def test_preferred_border_agent_id( + hass: HomeAssistant, hass_ws_client: WebSocketGenerator +) -> None: + """Test setting and getting the preferred border agent ID.""" + assert await async_setup_component(hass, DOMAIN, {}) + await hass.async_block_till_done() + + client = await hass_ws_client(hass) + + await client.send_json_auto_id({"type": "thread/get_preferred_border_agent_id"}) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == {"border_agent_id": None} + + await client.send_json_auto_id( + {"type": "thread/set_preferred_border_agent_id", "border_agent_id": "blah"} + ) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] is None + + await client.send_json_auto_id({"type": "thread/get_preferred_border_agent_id"}) + msg = await client.receive_json() + assert msg["success"] + assert msg["result"] == {"border_agent_id": "blah"} + + assert await dataset_store.async_get_preferred_border_agent_id(hass) == "blah" + + async def test_set_preferred_dataset( hass: HomeAssistant, hass_ws_client: WebSocketGenerator ) -> None: