Fix Home Assistant Yellow config entry data (#144948)

homeassistant/components/homeassistant_yellow/__init__.py

@@ -90,16 +90,17 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
                 minor_version=2,
             )
 
-        if config_entry.minor_version == 2:
-            # Add a `firmware_version` key
+        if config_entry.minor_version <= 3:
+            # Add a `firmware_version` key if it doesn't exist to handle entries created
+            # with minor version 1.3 where the firmware version was not set.
             hass.config_entries.async_update_entry(
                 config_entry,
                 data={
                     **config_entry.data,
-                    FIRMWARE_VERSION: None,
+                    FIRMWARE_VERSION: config_entry.data.get(FIRMWARE_VERSION),
                 },
                 version=1,
-                minor_version=3,
+                minor_version=4,
             )
 
         _LOGGER.debug(

homeassistant/components/homeassistant_yellow/config_flow.py

@@ -62,7 +62,7 @@ class HomeAssistantYellowConfigFlow(BaseFirmwareConfigFlow, domain=DOMAIN):
     """Handle a config flow for Home Assistant Yellow."""
 
     VERSION = 1
-    MINOR_VERSION = 3
+    MINOR_VERSION = 4
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Instantiate config flow."""
@@ -116,6 +116,11 @@ class HomeAssistantYellowConfigFlow(BaseFirmwareConfigFlow, domain=DOMAIN):
                     if self._probed_firmware_info is not None
                     else ApplicationType.EZSP
                 ).value,
+                FIRMWARE_VERSION: (
+                    self._probed_firmware_info.firmware_version
+                    if self._probed_firmware_info is not None
+                    else None
+                ),
             },
         )
 

tests/components/homeassistant_yellow/test_config_flow.py

@@ -101,12 +101,12 @@ async def test_config_flow(hass: HomeAssistant) -> None:
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["title"] == "Home Assistant Yellow"
-    assert result["data"] == {"firmware": "ezsp"}
+    assert result["data"] == {"firmware": "ezsp", "firmware_version": None}
     assert result["options"] == {}
     assert len(mock_setup_entry.mock_calls) == 1
 
     config_entry = hass.config_entries.async_entries(DOMAIN)[0]
-    assert config_entry.data == {"firmware": "ezsp"}
+    assert config_entry.data == {"firmware": "ezsp", "firmware_version": None}
     assert config_entry.options == {}
     assert config_entry.title == "Home Assistant Yellow"
 

tests/components/homeassistant_yellow/test_init.py

@@ -10,6 +10,9 @@ from homeassistant.components.homeassistant_hardware.util import (
     ApplicationType,
     FirmwareInfo,
 )
+from homeassistant.components.homeassistant_yellow.config_flow import (
+    HomeAssistantYellowConfigFlow,
+)
 from homeassistant.components.homeassistant_yellow.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
@@ -248,3 +251,71 @@ async def test_setup_entry_addon_info_fails(
 
     await hass.async_block_till_done()
     assert config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+@pytest.mark.parametrize(
+    ("start_version", "data", "migrated_data"),
+    [
+        (1, {}, {"firmware": "ezsp", "firmware_version": None}),
+        (2, {"firmware": "ezsp"}, {"firmware": "ezsp", "firmware_version": None}),
+        (
+            2,
+            {"firmware": "ezsp", "firmware_version": "123"},
+            {"firmware": "ezsp", "firmware_version": "123"},
+        ),
+        (3, {"firmware": "ezsp"}, {"firmware": "ezsp", "firmware_version": None}),
+        (
+            3,
+            {"firmware": "ezsp", "firmware_version": "123"},
+            {"firmware": "ezsp", "firmware_version": "123"},
+        ),
+    ],
+)
+async def test_migrate_entry(
+    hass: HomeAssistant,
+    start_version: int,
+    data: dict,
+    migrated_data: dict,
+) -> None:
+    """Test migration of a config entry."""
+    mock_integration(hass, MockModule("hassio"))
+    await async_setup_component(hass, HASSIO_DOMAIN, {})
+
+    # Setup the config entry
+    config_entry = MockConfigEntry(
+        data=data,
+        domain=DOMAIN,
+        options={},
+        title="Home Assistant Yellow",
+        version=1,
+        minor_version=start_version,
+    )
+    config_entry.add_to_hass(hass)
+
+    with (
+        patch(
+            "homeassistant.components.homeassistant_yellow.get_os_info",
+            return_value={"board": "yellow"},
+        ),
+        patch(
+            "homeassistant.components.onboarding.async_is_onboarded",
+            return_value=True,
+        ),
+        patch(
+            "homeassistant.components.homeassistant_yellow.guess_firmware_info",
+            return_value=FirmwareInfo(  # Nothing is setup
+                device="/dev/ttyAMA1",
+                firmware_version="1234",
+                firmware_type=ApplicationType.EZSP,
+                source="unknown",
+                owners=[],
+            ),
+        ),
+    ):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert config_entry.data == migrated_data
+    assert config_entry.options == {}
+    assert config_entry.minor_version == HomeAssistantYellowConfigFlow.MINOR_VERSION
+    assert config_entry.version == HomeAssistantYellowConfigFlow.VERSION