Store preferred border agent extended address for each thread dataset (#109065)

homeassistant/components/otbr/__init__.py

@@ -38,6 +38,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     try:
         border_agent_id = await otbrdata.get_border_agent_id()
         dataset_tlvs = await otbrdata.get_active_dataset_tlvs()
+        extended_address = await otbrdata.get_extended_address()
     except (
         HomeAssistantError,
         aiohttp.ClientError,
@@ -62,6 +63,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             DOMAIN,
             dataset_tlvs.hex(),
             preferred_border_agent_id=border_agent_id.hex(),
+            preferred_extended_address=extended_address.hex(),
         )
 
     entry.async_on_unload(entry.add_update_listener(async_reload_entry))

homeassistant/components/thread/dataset_store.py

@@ -23,7 +23,7 @@ BORDER_AGENT_DISCOVERY_TIMEOUT = 30
 DATA_STORE = "thread.datasets"
 STORAGE_KEY = "thread.datasets"
 STORAGE_VERSION_MAJOR = 1
-STORAGE_VERSION_MINOR = 3
+STORAGE_VERSION_MINOR = 4
 SAVE_DELAY = 10
 
 _LOGGER = logging.getLogger(__name__)
@@ -38,6 +38,7 @@ class DatasetEntry:
     """Dataset store entry."""
 
     preferred_border_agent_id: str | None
+    preferred_extended_address: str | None
     source: str
     tlv: str
 
@@ -79,6 +80,7 @@ class DatasetEntry:
             "created": self.created.isoformat(),
             "id": self.id,
             "preferred_border_agent_id": self.preferred_border_agent_id,
+            "preferred_extended_address": self.preferred_extended_address,
             "source": self.source,
             "tlv": self.tlv,
         }
@@ -104,6 +106,7 @@ class DatasetStoreStore(Store):
                         created=created,
                         id=dataset["id"],
                         preferred_border_agent_id=None,
+                        preferred_extended_address=None,
                         source=dataset["source"],
                         tlv=dataset["tlv"],
                     )
@@ -165,10 +168,14 @@ class DatasetStoreStore(Store):
                     "preferred_dataset": preferred_dataset,
                     "datasets": [dataset.to_json() for dataset in datasets.values()],
                 }
-            if old_minor_version < 3:
+            # Migration to version 1.3 removed, it added the ID of the preferred border
-                # Add border agent ID
+            # agent
+            if old_minor_version < 4:
+                # Add extended address of the preferred border agent and clear border
+                # agent ID
                 for dataset in data["datasets"]:
-                    dataset.setdefault("preferred_border_agent_id", None)
+                    dataset["preferred_border_agent_id"] = None
+                    dataset["preferred_extended_address"] = None
 
         return data
 
@@ -192,7 +199,11 @@ class DatasetStore:
 
     @callback
     def async_add(
-        self, source: str, tlv: str, preferred_border_agent_id: str | None
+        self,
+        source: str,
+        tlv: str,
+        preferred_border_agent_id: str | None,
+        preferred_extended_address: str | None,
     ) -> None:
         """Add dataset, does nothing if it already exists."""
         # Make sure the tlv is valid
@@ -206,16 +217,23 @@ class DatasetStore:
         ):
             raise HomeAssistantError("Invalid dataset")
 
+        # Don't allow setting preferred border agent ID without setting
+        # preferred extended address
+        if preferred_border_agent_id is not None and preferred_extended_address is None:
+            raise HomeAssistantError(
+                "Must set preferred extended address with preferred border agent ID"
+            )
+
         # Bail out if the dataset already exists
         entry: DatasetEntry | None
         for entry in self.datasets.values():
             if entry.dataset == dataset:
                 if (
-                    preferred_border_agent_id
+                    preferred_extended_address
-                    and entry.preferred_border_agent_id is None
+                    and entry.preferred_extended_address is None
                 ):
-                    self.async_set_preferred_border_agent_id(
+                    self.async_set_preferred_border_agent(
-                        entry.id, preferred_border_agent_id
+                        entry.id, preferred_border_agent_id, preferred_extended_address
                     )
                 return
 
@@ -262,14 +280,17 @@ class DatasetStore:
                 self.datasets[entry.id], tlv=tlv
             )
             self.async_schedule_save()
-            if preferred_border_agent_id and entry.preferred_border_agent_id is None:
+            if preferred_extended_address and entry.preferred_extended_address is None:
-                self.async_set_preferred_border_agent_id(
+                self.async_set_preferred_border_agent(
-                    entry.id, preferred_border_agent_id
+                    entry.id, preferred_border_agent_id, preferred_extended_address
                 )
             return
 
         entry = DatasetEntry(
-            preferred_border_agent_id=preferred_border_agent_id, source=source, tlv=tlv
+            preferred_border_agent_id=preferred_border_agent_id,
+            preferred_extended_address=preferred_extended_address,
+            source=source,
+            tlv=tlv,
         )
         self.datasets[entry.id] = entry
         self.async_schedule_save()
@@ -278,12 +299,12 @@ class DatasetStore:
         # no other router present. We only attempt this once.
         if (
             self._preferred_dataset is None
-            and preferred_border_agent_id
+            and preferred_extended_address
             and not self._set_preferred_dataset_task
         ):
             self._set_preferred_dataset_task = self.hass.async_create_task(
                 self._set_preferred_dataset_if_only_network(
-                    entry.id, preferred_border_agent_id
+                    entry.id, preferred_extended_address
                 )
             )
 
@@ -301,12 +322,21 @@ class DatasetStore:
         return self.datasets.get(dataset_id)
 
     @callback
-    def async_set_preferred_border_agent_id(
+    def async_set_preferred_border_agent(
-        self, dataset_id: str, border_agent_id: str
+        self, dataset_id: str, border_agent_id: str | None, extended_address: str
     ) -> None:
-        """Set preferred border agent id of a dataset."""
+        """Set preferred border agent id and extended address of a dataset."""
+        # Don't allow setting preferred border agent ID without setting
+        # preferred extended address
+        if border_agent_id is not None and extended_address is None:
+            raise HomeAssistantError(
+                "Must set preferred extended address with preferred border agent ID"
+            )
+
         self.datasets[dataset_id] = dataclasses.replace(
-            self.datasets[dataset_id], preferred_border_agent_id=border_agent_id
+            self.datasets[dataset_id],
+            preferred_border_agent_id=border_agent_id,
+            preferred_extended_address=extended_address,
         )
         self.async_schedule_save()
 
@@ -326,12 +356,12 @@ class DatasetStore:
         self.async_schedule_save()
 
     async def _set_preferred_dataset_if_only_network(
-        self, dataset_id: str, border_agent_id: str
+        self, dataset_id: str, extended_address: str | None
     ) -> None:
         """Set the preferred dataset, unless there are other routers present."""
         _LOGGER.debug(
             "_set_preferred_dataset_if_only_network called for router %s",
-            border_agent_id,
+            extended_address,
         )
 
         own_router_evt = Event()
@@ -342,8 +372,8 @@ class DatasetStore:
             key: str, data: discovery.ThreadRouterDiscoveryData
         ) -> None:
             """Handle router discovered."""
-            _LOGGER.debug("discovered router with id %s", data.border_agent_id)
+            _LOGGER.debug("discovered router with ext addr %s", data.extended_address)
-            if data.border_agent_id == border_agent_id:
+            if data.extended_address == extended_address:
                 own_router_evt.set()
                 return
 
@@ -395,6 +425,7 @@ class DatasetStore:
                     created=created,
                     id=dataset["id"],
                     preferred_border_agent_id=dataset["preferred_border_agent_id"],
+                    preferred_extended_address=dataset["preferred_extended_address"],
                     source=dataset["source"],
                     tlv=dataset["tlv"],
                 )
@@ -431,10 +462,11 @@ async def async_add_dataset(
     tlv: str,
     *,
     preferred_border_agent_id: str | None = None,
+    preferred_extended_address: str | None = None,
 ) -> None:
     """Add a dataset."""
     store = await async_get_store(hass)
-    store.async_add(source, tlv, preferred_border_agent_id)
+    store.async_add(source, tlv, preferred_border_agent_id, preferred_extended_address)
 
 
 async def async_get_dataset(hass: HomeAssistant, dataset_id: str) -> str | None:

homeassistant/components/thread/websocket_api.py

@@ -20,7 +20,7 @@ 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_set_preferred_border_agent_id)
+    websocket_api.async_register_command(hass, ws_set_preferred_border_agent)
     websocket_api.async_register_command(hass, ws_set_preferred_dataset)
 
 
@@ -54,20 +54,24 @@ async def ws_add_dataset(
 @websocket_api.require_admin
 @websocket_api.websocket_command(
     {
-        vol.Required("type"): "thread/set_preferred_border_agent_id",
+        vol.Required("type"): "thread/set_preferred_border_agent",
         vol.Required("dataset_id"): str,
-        vol.Required("border_agent_id"): str,
+        vol.Required("border_agent_id"): vol.Any(str, None),
+        vol.Required("extended_address"): str,
     }
 )
 @websocket_api.async_response
-async def ws_set_preferred_border_agent_id(
+async def ws_set_preferred_border_agent(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
 ) -> None:
-    """Set the preferred border agent ID."""
+    """Set the preferred border agent's border agent ID and extended address."""
     dataset_id = msg["dataset_id"]
     border_agent_id = msg["border_agent_id"]
+    extended_address = msg["extended_address"]
     store = await dataset_store.async_get_store(hass)
-    store.async_set_preferred_border_agent_id(dataset_id, border_agent_id)
+    store.async_set_preferred_border_agent(
+        dataset_id, border_agent_id, extended_address
+    )
     connection.send_result(msg["id"])
 
 
@@ -174,6 +178,7 @@ async def ws_list_datasets(
                 "pan_id": dataset.pan_id,
                 "preferred": dataset.id == preferred_dataset,
                 "preferred_border_agent_id": dataset.preferred_border_agent_id,
+                "preferred_extended_address": dataset.preferred_extended_address,
                 "source": dataset.source,
             }
         )

tests/components/otbr/__init__.py

@@ -27,8 +27,9 @@ DATASET_INSECURE_PASSPHRASE = bytes.fromhex(
     "0212340410445F2B5CA6F2A93A55CE570A70EFEECB0C0402A0F7F8"
 )
 
-TEST_BORDER_AGENT_ID = bytes.fromhex("230C6A1AC57F6F4BE262ACF32E5EF52C")
+TEST_BORDER_AGENT_EXTENDED_ADDRESS = bytes.fromhex("AEEB2F594B570BBF")
 
+TEST_BORDER_AGENT_ID = bytes.fromhex("230C6A1AC57F6F4BE262ACF32E5EF52C")
 
 ROUTER_DISCOVERY_HASS = {
     "type_": "_meshcop._udp.local.",

tests/components/otbr/conftest.py

@@ -10,6 +10,7 @@ from . import (
     CONFIG_ENTRY_DATA_MULTIPAN,
     CONFIG_ENTRY_DATA_THREAD,
     DATASET_CH16,
+    TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     TEST_BORDER_AGENT_ID,
 )
 
@@ -30,6 +31,9 @@ async def otbr_config_entry_multipan_fixture(hass):
         "python_otbr_api.OTBR.get_active_dataset_tlvs", return_value=DATASET_CH16
     ), patch(
         "python_otbr_api.OTBR.get_border_agent_id", return_value=TEST_BORDER_AGENT_ID
+    ), patch(
+        "python_otbr_api.OTBR.get_extended_address",
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     ), patch(
         "homeassistant.components.otbr.util.compute_pskc"
     ):  # Patch to speed up tests
@@ -50,6 +54,9 @@ async def otbr_config_entry_thread_fixture(hass):
         "python_otbr_api.OTBR.get_active_dataset_tlvs", return_value=DATASET_CH16
     ), patch(
         "python_otbr_api.OTBR.get_border_agent_id", return_value=TEST_BORDER_AGENT_ID
+    ), patch(
+        "python_otbr_api.OTBR.get_extended_address",
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     ), patch(
         "homeassistant.components.otbr.util.compute_pskc"
     ):  # Patch to speed up tests

tests/components/otbr/test_init.py

@@ -25,6 +25,7 @@ from . import (
     DATASET_INSECURE_NW_KEY,
     DATASET_INSECURE_PASSPHRASE,
     ROUTER_DISCOVERY_HASS,
+    TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     TEST_BORDER_AGENT_ID,
 )
 
@@ -66,6 +67,9 @@ async def test_import_dataset(hass: HomeAssistant, mock_async_zeroconf: None) ->
         "python_otbr_api.OTBR.get_active_dataset_tlvs", return_value=DATASET_CH16
     ), patch(
         "python_otbr_api.OTBR.get_border_agent_id", return_value=TEST_BORDER_AGENT_ID
+    ), patch(
+        "python_otbr_api.OTBR.get_extended_address",
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     ), patch(
         "homeassistant.components.thread.dataset_store.BORDER_AGENT_DISCOVERY_TIMEOUT",
         0.1,
@@ -97,6 +101,10 @@ async def test_import_dataset(hass: HomeAssistant, mock_async_zeroconf: None) ->
         list(dataset_store.datasets.values())[0].preferred_border_agent_id
         == TEST_BORDER_AGENT_ID.hex()
     )
+    assert (
+        list(dataset_store.datasets.values())[0].preferred_extended_address
+        == TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex()
+    )
     assert await thread.async_get_preferred_dataset(hass) == DATASET_CH16.hex()
     assert not issue_registry.async_get_issue(
         domain=otbr.DOMAIN, issue_id=f"insecure_thread_network_{config_entry.entry_id}"
@@ -130,13 +138,19 @@ async def test_import_share_radio_channel_collision(
         "python_otbr_api.OTBR.get_active_dataset_tlvs", return_value=DATASET_CH16
     ), patch(
         "python_otbr_api.OTBR.get_border_agent_id", return_value=TEST_BORDER_AGENT_ID
+    ), patch(
+        "python_otbr_api.OTBR.get_extended_address",
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     ), patch(
         "homeassistant.components.thread.dataset_store.DatasetStore.async_add"
     ) as mock_add:
         assert await hass.config_entries.async_setup(config_entry.entry_id)
 
     mock_add.assert_called_once_with(
-        otbr.DOMAIN, DATASET_CH16.hex(), TEST_BORDER_AGENT_ID.hex()
+        otbr.DOMAIN,
+        DATASET_CH16.hex(),
+        TEST_BORDER_AGENT_ID.hex(),
+        TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
     )
     assert issue_registry.async_get_issue(
         domain=otbr.DOMAIN,
@@ -167,13 +181,19 @@ async def test_import_share_radio_no_channel_collision(
         "python_otbr_api.OTBR.get_active_dataset_tlvs", return_value=dataset
     ), patch(
         "python_otbr_api.OTBR.get_border_agent_id", return_value=TEST_BORDER_AGENT_ID
+    ), patch(
+        "python_otbr_api.OTBR.get_extended_address",
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     ), patch(
         "homeassistant.components.thread.dataset_store.DatasetStore.async_add"
     ) as mock_add:
         assert await hass.config_entries.async_setup(config_entry.entry_id)
 
     mock_add.assert_called_once_with(
-        otbr.DOMAIN, dataset.hex(), TEST_BORDER_AGENT_ID.hex()
+        otbr.DOMAIN,
+        dataset.hex(),
+        TEST_BORDER_AGENT_ID.hex(),
+        TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
     )
     assert not issue_registry.async_get_issue(
         domain=otbr.DOMAIN,
@@ -202,13 +222,19 @@ async def test_import_insecure_dataset(hass: HomeAssistant, dataset: bytes) -> N
         "python_otbr_api.OTBR.get_active_dataset_tlvs", return_value=dataset
     ), patch(
         "python_otbr_api.OTBR.get_border_agent_id", return_value=TEST_BORDER_AGENT_ID
+    ), patch(
+        "python_otbr_api.OTBR.get_extended_address",
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     ), patch(
         "homeassistant.components.thread.dataset_store.DatasetStore.async_add"
     ) as mock_add:
         assert await hass.config_entries.async_setup(config_entry.entry_id)
 
     mock_add.assert_called_once_with(
-        otbr.DOMAIN, dataset.hex(), TEST_BORDER_AGENT_ID.hex()
+        otbr.DOMAIN,
+        dataset.hex(),
+        TEST_BORDER_AGENT_ID.hex(),
+        TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
     )
     assert issue_registry.async_get_issue(
         domain=otbr.DOMAIN, issue_id=f"insecure_thread_network_{config_entry.entry_id}"
@@ -268,6 +294,9 @@ async def test_config_entry_update(hass: HomeAssistant) -> None:
     mock_api = MagicMock()
     mock_api.get_active_dataset_tlvs = AsyncMock(return_value=None)
     mock_api.get_border_agent_id = AsyncMock(return_value=TEST_BORDER_AGENT_ID)
+    mock_api.get_extended_address = AsyncMock(
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS
+    )
     with patch("python_otbr_api.OTBR", return_value=mock_api) as mock_otrb_api:
         assert await hass.config_entries.async_setup(config_entry.entry_id)
 

tests/components/otbr/test_websocket_api.py

@@ -8,7 +8,13 @@ from homeassistant.components import otbr, thread
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
-from . import BASE_URL, DATASET_CH15, DATASET_CH16, TEST_BORDER_AGENT_ID
+from . import (
+    BASE_URL,
+    DATASET_CH15,
+    DATASET_CH16,
+    TEST_BORDER_AGENT_EXTENDED_ADDRESS,
+    TEST_BORDER_AGENT_ID,
+)
 
 from tests.test_util.aiohttp import AiohttpClientMocker
 from tests.typing import WebSocketGenerator
@@ -37,7 +43,7 @@ async def test_get_info(
         "python_otbr_api.OTBR.get_border_agent_id", return_value=TEST_BORDER_AGENT_ID
     ), patch(
         "python_otbr_api.OTBR.get_extended_address",
-        return_value=bytes.fromhex("4EF6C4F3FF750626"),
+        return_value=TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     ):
         await websocket_client.send_json_auto_id({"type": "otbr/info"})
         msg = await websocket_client.receive_json()
@@ -48,7 +54,7 @@ async def test_get_info(
         "active_dataset_tlvs": DATASET_CH16.hex().lower(),
         "channel": 16,
         "border_agent_id": TEST_BORDER_AGENT_ID.hex(),
-        "extended_address": "4EF6C4F3FF750626".lower(),
+        "extended_address": TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
     }
 
 
@@ -125,7 +131,7 @@ async def test_create_network(
     assert set_enabled_mock.mock_calls[0][1][0] is False
     assert set_enabled_mock.mock_calls[1][1][0] is True
     get_active_dataset_tlvs_mock.assert_called_once()
-    mock_add.assert_called_once_with(otbr.DOMAIN, DATASET_CH16.hex(), None)
+    mock_add.assert_called_once_with(otbr.DOMAIN, DATASET_CH16.hex(), None, None)
 
 
 async def test_create_network_no_entry(

tests/components/thread/__init__.py

@@ -18,6 +18,8 @@ DATASET_3 = (
     "0212340410445F2B5CA6F2A93A55CE570A70EFEECB0C0402A0F7F8"
 )
 
+TEST_BORDER_AGENT_EXTENDED_ADDRESS = bytes.fromhex("AEEB2F594B570BBF")
+
 TEST_BORDER_AGENT_ID = bytes.fromhex("230C6A1AC57F6F4BE262ACF32E5EF52C")
 
 ROUTER_DISCOVERY_GOOGLE_1 = {

tests/components/thread/test_dataset_store.py

@@ -17,6 +17,7 @@ from . import (
     DATASET_3,
     ROUTER_DISCOVERY_GOOGLE_1,
     ROUTER_DISCOVERY_HASS,
+    TEST_BORDER_AGENT_EXTENDED_ADDRESS,
     TEST_BORDER_AGENT_ID,
 )
 
@@ -269,7 +270,7 @@ async def test_load_datasets(hass: HomeAssistant) -> None:
 
     store1 = await dataset_store.async_get_store(hass)
     for dataset in datasets:
-        store1.async_add(dataset["source"], dataset["tlv"], None)
+        store1.async_add(dataset["source"], dataset["tlv"], None, None)
     assert len(store1.datasets) == 3
     dataset_id = list(store1.datasets.values())[0].id
     store1.preferred_dataset = dataset_id
@@ -321,6 +322,7 @@ async def test_loading_datasets_from_storage(
                     "created": "2023-02-02T09:41:13.746514+00:00",
                     "id": "id1",
                     "preferred_border_agent_id": "230C6A1AC57F6F4BE262ACF32E5EF52C",
+                    "preferred_extended_address": "AEEB2F594B570BBF",
                     "source": "source_1",
                     "tlv": DATASET_1,
                 },
@@ -328,6 +330,7 @@ async def test_loading_datasets_from_storage(
                     "created": "2023-02-02T09:41:13.746514+00:00",
                     "id": "id2",
                     "preferred_border_agent_id": None,
+                    "preferred_extended_address": "AEEB2F594B570BBF",
                     "source": "source_2",
                     "tlv": DATASET_2,
                 },
@@ -335,6 +338,7 @@ async def test_loading_datasets_from_storage(
                     "created": "2023-02-02T09:41:13.746514+00:00",
                     "id": "id3",
                     "preferred_border_agent_id": None,
+                    "preferred_extended_address": None,
                     "source": "source_3",
                     "tlv": DATASET_3,
                 },
@@ -556,42 +560,148 @@ async def test_migrate_set_default_border_agent_id(
 
     store = await dataset_store.async_get_store(hass)
     assert store.datasets[store._preferred_dataset].preferred_border_agent_id is None
+    assert store.datasets[store._preferred_dataset].preferred_extended_address is None
 
 
 async def test_set_preferred_border_agent_id(hass: HomeAssistant) -> None:
     """Test set the preferred border agent ID of a dataset."""
     assert await dataset_store.async_get_preferred_dataset(hass) is None
 
+    with pytest.raises(HomeAssistantError):
         await dataset_store.async_add_dataset(
             hass, "source", DATASET_3, preferred_border_agent_id="blah"
         )
 
     store = await dataset_store.async_get_store(hass)
-    assert len(store.datasets) == 1
+    assert len(store.datasets) == 0
-    assert list(store.datasets.values())[0].preferred_border_agent_id == "blah"
 
+    with pytest.raises(HomeAssistantError):
         await dataset_store.async_add_dataset(
             hass, "source", DATASET_3, preferred_border_agent_id="bleh"
         )
+    assert len(store.datasets) == 0
+
+    await dataset_store.async_add_dataset(hass, "source", DATASET_2)
+    assert len(store.datasets) == 1
+    assert list(store.datasets.values())[0].preferred_border_agent_id is None
+
+    with pytest.raises(HomeAssistantError):
+        await dataset_store.async_add_dataset(
+            hass, "source", DATASET_2, preferred_border_agent_id="blah"
+        )
+    assert list(store.datasets.values())[0].preferred_border_agent_id is None
+
+    store = await dataset_store.async_get_store(hass)
+    dataset_id = list(store.datasets.values())[0].id
+    with pytest.raises(HomeAssistantError):
+        await store.async_set_preferred_border_agent(dataset_id, "blah", None)
+    assert list(store.datasets.values())[0].preferred_border_agent_id is None
+
+    await dataset_store.async_add_dataset(hass, "source", DATASET_1)
+    assert len(store.datasets) == 2
+    assert list(store.datasets.values())[1].preferred_border_agent_id is None
+
+    with pytest.raises(HomeAssistantError):
+        await dataset_store.async_add_dataset(
+            hass, "source", DATASET_1_LARGER_TIMESTAMP, preferred_border_agent_id="blah"
+        )
+    assert list(store.datasets.values())[1].preferred_border_agent_id is None
+
+
+async def test_set_preferred_border_agent_id_and_extended_address(
+    hass: HomeAssistant,
+) -> None:
+    """Test set the preferred border agent ID and extended address of a dataset."""
+    assert await dataset_store.async_get_preferred_dataset(hass) is None
+
+    await dataset_store.async_add_dataset(
+        hass,
+        "source",
+        DATASET_3,
+        preferred_border_agent_id="blah",
+        preferred_extended_address="bleh",
+    )
+
+    store = await dataset_store.async_get_store(hass)
+    assert len(store.datasets) == 1
     assert list(store.datasets.values())[0].preferred_border_agent_id == "blah"
+    assert list(store.datasets.values())[0].preferred_extended_address == "bleh"
+
+    await dataset_store.async_add_dataset(
+        hass,
+        "source",
+        DATASET_3,
+        preferred_border_agent_id="bleh",
+        preferred_extended_address="bleh",
+    )
+    assert list(store.datasets.values())[0].preferred_border_agent_id == "blah"
+    assert list(store.datasets.values())[0].preferred_extended_address == "bleh"
 
     await dataset_store.async_add_dataset(hass, "source", DATASET_2)
     assert len(store.datasets) == 2
     assert list(store.datasets.values())[1].preferred_border_agent_id is None
+    assert list(store.datasets.values())[1].preferred_extended_address is None
 
     await dataset_store.async_add_dataset(
-        hass, "source", DATASET_2, preferred_border_agent_id="blah"
+        hass,
+        "source",
+        DATASET_2,
+        preferred_border_agent_id="blah",
+        preferred_extended_address="bleh",
     )
     assert list(store.datasets.values())[1].preferred_border_agent_id == "blah"
+    assert list(store.datasets.values())[1].preferred_extended_address == "bleh"
 
     await dataset_store.async_add_dataset(hass, "source", DATASET_1)
     assert len(store.datasets) == 3
     assert list(store.datasets.values())[2].preferred_border_agent_id is None
+    assert list(store.datasets.values())[2].preferred_extended_address is None
 
     await dataset_store.async_add_dataset(
-        hass, "source", DATASET_1_LARGER_TIMESTAMP, preferred_border_agent_id="blah"
+        hass,
+        "source",
+        DATASET_1_LARGER_TIMESTAMP,
+        preferred_border_agent_id="blah",
+        preferred_extended_address="bleh",
     )
-    assert list(store.datasets.values())[1].preferred_border_agent_id == "blah"
+    assert list(store.datasets.values())[2].preferred_border_agent_id == "blah"
+    assert list(store.datasets.values())[2].preferred_extended_address == "bleh"
+
+
+async def test_set_preferred_extended_address(hass: HomeAssistant) -> None:
+    """Test set the preferred extended address of a dataset."""
+    assert await dataset_store.async_get_preferred_dataset(hass) is None
+
+    await dataset_store.async_add_dataset(
+        hass, "source", DATASET_3, preferred_extended_address="blah"
+    )
+
+    store = await dataset_store.async_get_store(hass)
+    assert len(store.datasets) == 1
+    assert list(store.datasets.values())[0].preferred_extended_address == "blah"
+
+    await dataset_store.async_add_dataset(
+        hass, "source", DATASET_3, preferred_extended_address="bleh"
+    )
+    assert list(store.datasets.values())[0].preferred_extended_address == "blah"
+
+    await dataset_store.async_add_dataset(hass, "source", DATASET_2)
+    assert len(store.datasets) == 2
+    assert list(store.datasets.values())[1].preferred_extended_address is None
+
+    await dataset_store.async_add_dataset(
+        hass, "source", DATASET_2, preferred_extended_address="blah"
+    )
+    assert list(store.datasets.values())[1].preferred_extended_address == "blah"
+
+    await dataset_store.async_add_dataset(hass, "source", DATASET_1)
+    assert len(store.datasets) == 3
+    assert list(store.datasets.values())[2].preferred_extended_address is None
+
+    await dataset_store.async_add_dataset(
+        hass, "source", DATASET_1_LARGER_TIMESTAMP, preferred_extended_address="blah"
+    )
+    assert list(store.datasets.values())[2].preferred_extended_address == "blah"
 
 
 async def test_automatically_set_preferred_dataset(
@@ -624,6 +734,7 @@ async def test_automatically_set_preferred_dataset(
             "source",
             DATASET_1,
             preferred_border_agent_id=TEST_BORDER_AGENT_ID.hex(),
+            preferred_extended_address=TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
         )
 
         # Wait for discovery to start
@@ -651,6 +762,10 @@ async def test_automatically_set_preferred_dataset(
         list(store.datasets.values())[0].preferred_border_agent_id
         == TEST_BORDER_AGENT_ID.hex()
     )
+    assert (
+        list(store.datasets.values())[0].preferred_extended_address
+        == TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex()
+    )
     assert await dataset_store.async_get_preferred_dataset(hass) == DATASET_1
 
 
@@ -687,6 +802,7 @@ async def test_automatically_set_preferred_dataset_own_and_other_router(
             "source",
             DATASET_1,
             preferred_border_agent_id=TEST_BORDER_AGENT_ID.hex(),
+            preferred_extended_address=TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
         )
 
         # Wait for discovery to start
@@ -725,6 +841,10 @@ async def test_automatically_set_preferred_dataset_own_and_other_router(
         list(store.datasets.values())[0].preferred_border_agent_id
         == TEST_BORDER_AGENT_ID.hex()
     )
+    assert (
+        list(store.datasets.values())[0].preferred_extended_address
+        == TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex()
+    )
     assert await dataset_store.async_get_preferred_dataset(hass) is None
 
 
@@ -761,6 +881,7 @@ async def test_automatically_set_preferred_dataset_other_router(
             "source",
             DATASET_1,
             preferred_border_agent_id=TEST_BORDER_AGENT_ID.hex(),
+            preferred_extended_address=TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
         )
 
         # Wait for discovery to start
@@ -788,6 +909,10 @@ async def test_automatically_set_preferred_dataset_other_router(
         list(store.datasets.values())[0].preferred_border_agent_id
         == TEST_BORDER_AGENT_ID.hex()
     )
+    assert (
+        list(store.datasets.values())[0].preferred_extended_address
+        == TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex()
+    )
     assert await dataset_store.async_get_preferred_dataset(hass) is None
 
 
@@ -824,6 +949,7 @@ async def test_automatically_set_preferred_dataset_no_router(
             "source",
             DATASET_1,
             preferred_border_agent_id=TEST_BORDER_AGENT_ID.hex(),
+            preferred_extended_address=TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex(),
         )
 
         # Wait for discovery to start
@@ -840,4 +966,8 @@ async def test_automatically_set_preferred_dataset_no_router(
         list(store.datasets.values())[0].preferred_border_agent_id
         == TEST_BORDER_AGENT_ID.hex()
     )
+    assert (
+        list(store.datasets.values())[0].preferred_extended_address
+        == TEST_BORDER_AGENT_EXTENDED_ADDRESS.hex()
+    )
     assert await dataset_store.async_get_preferred_dataset(hass) is None

tests/components/thread/test_websocket_api.py

@@ -175,6 +175,7 @@ async def test_list_get_dataset(
                 "pan_id": "1234",
                 "preferred": True,
                 "preferred_border_agent_id": None,
+                "preferred_extended_address": None,
                 "source": "Google",
             },
             {
@@ -186,6 +187,7 @@ async def test_list_get_dataset(
                 "pan_id": "1234",
                 "preferred": False,
                 "preferred_border_agent_id": None,
+                "preferred_extended_address": None,
                 "source": "Multipan",
             },
             {
@@ -197,6 +199,7 @@ async def test_list_get_dataset(
                 "pan_id": "1234",
                 "preferred": False,
                 "preferred_border_agent_id": None,
+                "preferred_extended_address": None,
                 "source": "🎅",
             },
         ]
@@ -217,7 +220,7 @@ async def test_list_get_dataset(
     assert msg["error"] == {"code": "not_found", "message": "unknown dataset"}
 
 
-async def test_set_preferred_border_agent_id(
+async def test_set_preferred_border_agent(
     hass: HomeAssistant, hass_ws_client: WebSocketGenerator
 ) -> None:
     """Test setting the preferred border agent ID."""
@@ -239,12 +242,14 @@ async def test_set_preferred_border_agent_id(
     datasets = msg["result"]["datasets"]
     dataset_id = datasets[0]["dataset_id"]
     assert datasets[0]["preferred_border_agent_id"] is None
+    assert datasets[0]["preferred_extended_address"] is None
 
     await client.send_json_auto_id(
         {
-            "type": "thread/set_preferred_border_agent_id",
+            "type": "thread/set_preferred_border_agent",
             "dataset_id": dataset_id,
             "border_agent_id": "blah",
+            "extended_address": "bleh",
         }
     )
     msg = await client.receive_json()
@@ -256,6 +261,7 @@ async def test_set_preferred_border_agent_id(
     assert msg["success"]
     datasets = msg["result"]["datasets"]
     assert datasets[0]["preferred_border_agent_id"] == "blah"
+    assert datasets[0]["preferred_extended_address"] == "bleh"
 
 
 async def test_set_preferred_dataset(