From 7dcf2e94b4401c4ace7168cca59b8d07acbebf84 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 12 Feb 2024 14:24:56 -0600
Subject: [PATCH] Migrate zha to use async_update_entry to alter config entries
 (#110402)

---
 homeassistant/components/zha/__init__.py |  9 +++------
 tests/components/zha/test_config_flow.py |  5 +++--
 tests/components/zha/test_gateway.py     | 10 ++++++----
 tests/components/zha/test_init.py        | 24 ++++++++++++++----------
 4 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/homeassistant/components/zha/__init__.py b/homeassistant/components/zha/__init__.py
index 1eb3369c1be..c6666911ef9 100644
--- a/homeassistant/components/zha/__init__.py
+++ b/homeassistant/components/zha/__init__.py
@@ -272,8 +272,7 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         if data[CONF_RADIO_TYPE] != RadioType.deconz and baudrate in BAUD_RATES:
             data[CONF_DEVICE][CONF_BAUDRATE] = baudrate
 
-        config_entry.version = 2
-        hass.config_entries.async_update_entry(config_entry, data=data)
+        hass.config_entries.async_update_entry(config_entry, data=data, version=2)
 
     if config_entry.version == 2:
         data = {**config_entry.data}
@@ -281,8 +280,7 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         if data[CONF_RADIO_TYPE] == "ti_cc":
             data[CONF_RADIO_TYPE] = "znp"
 
-        config_entry.version = 3
-        hass.config_entries.async_update_entry(config_entry, data=data)
+        hass.config_entries.async_update_entry(config_entry, data=data, version=3)
 
     if config_entry.version == 3:
         data = {**config_entry.data}
@@ -299,8 +297,7 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         if not data[CONF_DEVICE].get(CONF_FLOW_CONTROL):
             data[CONF_DEVICE][CONF_FLOW_CONTROL] = None
 
-        config_entry.version = 4
-        hass.config_entries.async_update_entry(config_entry, data=data)
+        hass.config_entries.async_update_entry(config_entry, data=data, version=4)
 
     _LOGGER.info("Migration to version %s successful", config_entry.version)
     return True
diff --git a/tests/components/zha/test_config_flow.py b/tests/components/zha/test_config_flow.py
index 29a996d4477..e95c39e5f82 100644
--- a/tests/components/zha/test_config_flow.py
+++ b/tests/components/zha/test_config_flow.py
@@ -1953,9 +1953,10 @@ async def test_migration_ti_cc_to_znp(
     old_type: str, new_type: str, hass: HomeAssistant, config_entry: MockConfigEntry
 ) -> None:
     """Test zigpy-cc to zigpy-znp config migration."""
-    config_entry.data = {**config_entry.data, CONF_RADIO_TYPE: old_type}
-    config_entry.version = 2
     config_entry.add_to_hass(hass)
+    hass.config_entries.async_update_entry(
+        config_entry, data={**config_entry.data, CONF_RADIO_TYPE: old_type}, version=2
+    )
 
     with patch("homeassistant.components.zha.async_setup_entry", return_value=True):
         await hass.config_entries.async_setup(config_entry.entry_id)
diff --git a/tests/components/zha/test_gateway.py b/tests/components/zha/test_gateway.py
index ec467df9e52..e4d8d2a5d65 100644
--- a/tests/components/zha/test_gateway.py
+++ b/tests/components/zha/test_gateway.py
@@ -250,9 +250,10 @@ async def test_gateway_initialize_bellows_thread(
     config_entry: MockConfigEntry,
 ) -> None:
     """Test ZHA disabling the UART thread when connecting to a TCP coordinator."""
-    config_entry.data = dict(config_entry.data)
-    config_entry.data["device"]["path"] = device_path
+    data = dict(config_entry.data)
+    data["device"]["path"] = device_path
     config_entry.add_to_hass(hass)
+    hass.config_entries.async_update_entry(config_entry, data=data)
 
     zha_gateway = ZHAGateway(hass, {"zigpy_config": config_override}, config_entry)
 
@@ -285,9 +286,10 @@ async def test_gateway_force_multi_pan_channel(
     config_entry: MockConfigEntry,
 ) -> None:
     """Test ZHA disabling the UART thread when connecting to a TCP coordinator."""
-    config_entry.data = dict(config_entry.data)
-    config_entry.data["device"]["path"] = device_path
+    data = dict(config_entry.data)
+    data["device"]["path"] = device_path
     config_entry.add_to_hass(hass)
+    hass.config_entries.async_update_entry(config_entry, data=data)
 
     zha_gateway = ZHAGateway(hass, {"zigpy_config": config_override}, config_entry)
 
diff --git a/tests/components/zha/test_init.py b/tests/components/zha/test_init.py
index 9c9ffbb3ba1..5887fa2d8bc 100644
--- a/tests/components/zha/test_init.py
+++ b/tests/components/zha/test_init.py
@@ -199,17 +199,21 @@ async def test_migration_baudrate_and_flow_control(
     config_entry: MockConfigEntry,
 ) -> None:
     """Test baudrate and flow control migration."""
-    config_entry.data = {
-        **config_entry.data,
-        CONF_RADIO_TYPE: radio_type,
-        CONF_DEVICE: {
-            CONF_BAUDRATE: old_baudrate,
-            CONF_FLOW_CONTROL: old_flow_control,
-            CONF_DEVICE_PATH: "/dev/null",
-        },
-    }
-    config_entry.version = 3
+
     config_entry.add_to_hass(hass)
+    hass.config_entries.async_update_entry(
+        config_entry,
+        data={
+            **config_entry.data,
+            CONF_RADIO_TYPE: radio_type,
+            CONF_DEVICE: {
+                CONF_BAUDRATE: old_baudrate,
+                CONF_FLOW_CONTROL: old_flow_control,
+                CONF_DEVICE_PATH: "/dev/null",
+            },
+        },
+        version=3,
+    )
 
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()