Store preferred border agent ID in thread dataset store (#98375)

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)

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)

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(
     {

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"

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: