From 6895081595a19f4721ad7610e600fc61b2037706 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Mon, 6 Sep 2021 16:57:59 -0400
Subject: [PATCH] Use `async_update_entry` in config unit test instead of
 modifying `data` (#55855)

---
 homeassistant/components/zha/__init__.py | 4 ++--
 tests/components/zha/test_config_flow.py | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/zha/__init__.py b/homeassistant/components/zha/__init__.py
index b50adf15020..d6578be775f 100644
--- a/homeassistant/components/zha/__init__.py
+++ b/homeassistant/components/zha/__init__.py
@@ -184,7 +184,7 @@ async def async_migrate_entry(
             data[CONF_DEVICE][CONF_BAUDRATE] = baudrate
 
         config_entry.version = 2
-        config_entry.data = data
+        hass.config_entries.async_update_entry(config_entry, data=data)
 
     if config_entry.version == 2:
         data = {**config_entry.data}
@@ -193,7 +193,7 @@ async def async_migrate_entry(
             data[CONF_RADIO_TYPE] = "znp"
 
         config_entry.version = 3
-        config_entry.data = data
+        hass.config_entries.async_update_entry(config_entry, data=data)
 
     _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 b8a951a57f6..5aef30c854d 100644
--- a/tests/components/zha/test_config_flow.py
+++ b/tests/components/zha/test_config_flow.py
@@ -13,7 +13,7 @@ from homeassistant.components.ssdp import (
     ATTR_UPNP_MANUFACTURER_URL,
     ATTR_UPNP_SERIAL,
 )
-from homeassistant.components.zha import async_migrate_entry, config_flow
+from homeassistant.components.zha import config_flow
 from homeassistant.components.zha.core.const import (
     CONF_BAUDRATE,
     CONF_FLOWCONTROL,
@@ -633,8 +633,6 @@ async def test_user_port_config(probe_mock, hass):
 )
 async def test_migration_ti_cc_to_znp(old_type, new_type, hass, config_entry):
     """Test zigpy-cc to zigpy-znp config migration."""
-    await setup.async_setup_component(hass, "persistent_notification", {})
-
     config_entry = MockConfigEntry(
         domain=DOMAIN,
         unique_id=old_type + new_type,
@@ -651,7 +649,9 @@ async def test_migration_ti_cc_to_znp(old_type, new_type, hass, config_entry):
     config_entry.version = 2
     config_entry.add_to_hass(hass)
 
-    await async_migrate_entry(hass, config_entry)
+    with patch("homeassistant.components.zha.async_setup_entry", return_value=True):
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
 
     assert config_entry.version > 2
     assert config_entry.data[CONF_RADIO_TYPE] == new_type