From a213ef24757e95413e935da169e926dec29c20eb Mon Sep 17 00:00:00 2001
From: Bram Kragten <mail@bramkragten.nl>
Date: Tue, 14 Mar 2023 15:27:31 +0100
Subject: [PATCH] Add websocket command to set preferred thread dataset
 (#89700)

---
 .../components/thread/websocket_api.py        | 26 ++++++++++
 tests/components/thread/test_websocket_api.py | 50 +++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/homeassistant/components/thread/websocket_api.py b/homeassistant/components/thread/websocket_api.py
index 053ec69a0fa..9f9bc3455a8 100644
--- a/homeassistant/components/thread/websocket_api.py
+++ b/homeassistant/components/thread/websocket_api.py
@@ -20,6 +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_dataset)
 
 
 @websocket_api.require_admin
@@ -49,6 +50,31 @@ async def ws_add_dataset(
     connection.send_result(msg["id"])
 
 
+@websocket_api.require_admin
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "thread/set_preferred_dataset",
+        vol.Required("dataset_id"): str,
+    }
+)
+@websocket_api.async_response
+async def ws_set_preferred_dataset(
+    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Add a thread dataset."""
+    dataset_id = msg["dataset_id"]
+
+    store = await dataset_store.async_get_store(hass)
+    if not (store.async_get(dataset_id)):
+        connection.send_error(
+            msg["id"], websocket_api.const.ERR_NOT_FOUND, "unknown dataset"
+        )
+        return
+
+    store.preferred_dataset = dataset_id
+    connection.send_result(msg["id"])
+
+
 @websocket_api.require_admin
 @websocket_api.websocket_command(
     {
diff --git a/tests/components/thread/test_websocket_api.py b/tests/components/thread/test_websocket_api.py
index f8f09b0b8cc..c2e9e5f5934 100644
--- a/tests/components/thread/test_websocket_api.py
+++ b/tests/components/thread/test_websocket_api.py
@@ -197,6 +197,56 @@ async def test_list_get_dataset(
     assert msg["error"] == {"code": "not_found", "message": "unknown dataset"}
 
 
+async def test_set_preferred_dataset(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+) -> None:
+    """Test we set a dataset as default."""
+    assert await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    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 == "🎅":
+            dataset_3 = dataset
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json(
+        {"id": 1, "type": "thread/set_preferred_dataset", "dataset_id": dataset_3.id}
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None
+
+    store = await dataset_store.async_get_store(hass)
+    assert store.preferred_dataset == dataset_3.id
+
+
+async def test_set_preferred_dataset_wrong_id(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+) -> None:
+    """Test we set a dataset as default."""
+    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/set_preferred_dataset", "dataset_id": "don_t_exist"}
+    )
+    msg = await client.receive_json()
+    assert msg["error"]["code"] == "not_found"
+
+
 async def test_discover_routers(
     hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_async_zeroconf: None
 ) -> None: