Migrate zha to use async_update_entry to alter config entries (#110402)

This commit is contained in:
J. Nick Koston 2024-02-12 14:24:56 -06:00 committed by GitHub
parent c3d9192384
commit 7dcf2e94b4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 26 additions and 22 deletions

View file

@ -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: if data[CONF_RADIO_TYPE] != RadioType.deconz and baudrate in BAUD_RATES:
data[CONF_DEVICE][CONF_BAUDRATE] = baudrate data[CONF_DEVICE][CONF_BAUDRATE] = baudrate
config_entry.version = 2 hass.config_entries.async_update_entry(config_entry, data=data, version=2)
hass.config_entries.async_update_entry(config_entry, data=data)
if config_entry.version == 2: if config_entry.version == 2:
data = {**config_entry.data} 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": if data[CONF_RADIO_TYPE] == "ti_cc":
data[CONF_RADIO_TYPE] = "znp" data[CONF_RADIO_TYPE] = "znp"
config_entry.version = 3 hass.config_entries.async_update_entry(config_entry, data=data, version=3)
hass.config_entries.async_update_entry(config_entry, data=data)
if config_entry.version == 3: if config_entry.version == 3:
data = {**config_entry.data} 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): if not data[CONF_DEVICE].get(CONF_FLOW_CONTROL):
data[CONF_DEVICE][CONF_FLOW_CONTROL] = None data[CONF_DEVICE][CONF_FLOW_CONTROL] = None
config_entry.version = 4 hass.config_entries.async_update_entry(config_entry, data=data, version=4)
hass.config_entries.async_update_entry(config_entry, data=data)
_LOGGER.info("Migration to version %s successful", config_entry.version) _LOGGER.info("Migration to version %s successful", config_entry.version)
return True return True

View file

@ -1953,9 +1953,10 @@ async def test_migration_ti_cc_to_znp(
old_type: str, new_type: str, hass: HomeAssistant, config_entry: MockConfigEntry old_type: str, new_type: str, hass: HomeAssistant, config_entry: MockConfigEntry
) -> None: ) -> None:
"""Test zigpy-cc to zigpy-znp config migration.""" """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) 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): with patch("homeassistant.components.zha.async_setup_entry", return_value=True):
await hass.config_entries.async_setup(config_entry.entry_id) await hass.config_entries.async_setup(config_entry.entry_id)

View file

@ -250,9 +250,10 @@ async def test_gateway_initialize_bellows_thread(
config_entry: MockConfigEntry, config_entry: MockConfigEntry,
) -> None: ) -> None:
"""Test ZHA disabling the UART thread when connecting to a TCP coordinator.""" """Test ZHA disabling the UART thread when connecting to a TCP coordinator."""
config_entry.data = dict(config_entry.data) data = dict(config_entry.data)
config_entry.data["device"]["path"] = device_path data["device"]["path"] = device_path
config_entry.add_to_hass(hass) 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) 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, config_entry: MockConfigEntry,
) -> None: ) -> None:
"""Test ZHA disabling the UART thread when connecting to a TCP coordinator.""" """Test ZHA disabling the UART thread when connecting to a TCP coordinator."""
config_entry.data = dict(config_entry.data) data = dict(config_entry.data)
config_entry.data["device"]["path"] = device_path data["device"]["path"] = device_path
config_entry.add_to_hass(hass) 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) zha_gateway = ZHAGateway(hass, {"zigpy_config": config_override}, config_entry)

View file

@ -199,7 +199,11 @@ async def test_migration_baudrate_and_flow_control(
config_entry: MockConfigEntry, config_entry: MockConfigEntry,
) -> None: ) -> None:
"""Test baudrate and flow control migration.""" """Test baudrate and flow control migration."""
config_entry.data = {
config_entry.add_to_hass(hass)
hass.config_entries.async_update_entry(
config_entry,
data={
**config_entry.data, **config_entry.data,
CONF_RADIO_TYPE: radio_type, CONF_RADIO_TYPE: radio_type,
CONF_DEVICE: { CONF_DEVICE: {
@ -207,9 +211,9 @@ async def test_migration_baudrate_and_flow_control(
CONF_FLOW_CONTROL: old_flow_control, CONF_FLOW_CONTROL: old_flow_control,
CONF_DEVICE_PATH: "/dev/null", CONF_DEVICE_PATH: "/dev/null",
}, },
} },
config_entry.version = 3 version=3,
config_entry.add_to_hass(hass) )
await hass.config_entries.async_setup(config_entry.entry_id) await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done() await hass.async_block_till_done()