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

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) == 0
+
+    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"
+        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"
+        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(