From cc3783f9998f3d2210d39f00d3644fe4b3b41f3c Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 14 Feb 2024 16:10:07 -0600
Subject: [PATCH] Fix misaligned mac formatting in goalzero (#110574)

The DHCP mac address was being set as aabbcceeddff but the user step
was setting it as aa:bb:cc:dd:ee:ff
---
 homeassistant/components/goalzero/__init__.py | 14 +++++++++++++
 .../components/goalzero/config_flow.py        |  2 +-
 tests/components/goalzero/__init__.py         |  2 +-
 tests/components/goalzero/test_init.py        | 20 +++++++++++++++++++
 4 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/goalzero/__init__.py b/homeassistant/components/goalzero/__init__.py
index f32cad5a488..7a7000ba780 100644
--- a/homeassistant/components/goalzero/__init__.py
+++ b/homeassistant/components/goalzero/__init__.py
@@ -1,6 +1,8 @@
 """The Goal Zero Yeti integration."""
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from goalzero import Yeti, exceptions
 
 from homeassistant.config_entries import ConfigEntry
@@ -8,6 +10,7 @@ from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.device_registry import format_mac
 
 from .const import DOMAIN
 from .coordinator import GoalZeroDataUpdateCoordinator
@@ -17,6 +20,17 @@ PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH]
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Goal Zero Yeti from a config entry."""
+
+    mac = entry.unique_id
+
+    if TYPE_CHECKING:
+        assert mac is not None
+
+    if (formatted_mac := format_mac(mac)) != mac:
+        # The DHCP discovery path did not format the MAC address
+        # so we need to update the config entry if it's different
+        hass.config_entries.async_update_entry(entry, unique_id=formatted_mac)
+
     api = Yeti(entry.data[CONF_HOST], async_get_clientsession(hass))
     try:
         await api.init_connect()
diff --git a/homeassistant/components/goalzero/config_flow.py b/homeassistant/components/goalzero/config_flow.py
index 2d8c0c848c9..2312b6bd183 100644
--- a/homeassistant/components/goalzero/config_flow.py
+++ b/homeassistant/components/goalzero/config_flow.py
@@ -32,7 +32,7 @@ class GoalZeroFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         """Handle dhcp discovery."""
         self.ip_address = discovery_info.ip
 
-        await self.async_set_unique_id(discovery_info.macaddress)
+        await self.async_set_unique_id(format_mac(discovery_info.macaddress))
         self._abort_if_unique_id_configured(updates={CONF_HOST: self.ip_address})
         self._async_abort_entries_match({CONF_HOST: self.ip_address})
 
diff --git a/tests/components/goalzero/__init__.py b/tests/components/goalzero/__init__.py
index 18ac689017f..cde970f67a3 100644
--- a/tests/components/goalzero/__init__.py
+++ b/tests/components/goalzero/__init__.py
@@ -21,7 +21,7 @@ CONF_DATA = {
 
 CONF_DHCP_FLOW = dhcp.DhcpServiceInfo(
     ip=HOST,
-    macaddress=format_mac("AA:BB:CC:DD:EE:FF"),
+    macaddress=format_mac("AA:BB:CC:DD:EE:FF").replace(":", ""),
     hostname="yeti",
 )
 
diff --git a/tests/components/goalzero/test_init.py b/tests/components/goalzero/test_init.py
index 539da7e91a2..3a277d4cb53 100644
--- a/tests/components/goalzero/test_init.py
+++ b/tests/components/goalzero/test_init.py
@@ -35,6 +35,26 @@ async def test_setup_config_and_unload(hass: HomeAssistant) -> None:
     assert not hass.data.get(DOMAIN)
 
 
+async def test_setup_config_entry_incorrectly_formatted_mac(
+    hass: HomeAssistant,
+) -> None:
+    """Test the mac address formatting is corrected."""
+    entry = create_entry(hass)
+    hass.config_entries.async_update_entry(entry, unique_id="AABBCCDDEEFF")
+    mocked_yeti = await create_mocked_yeti()
+    with patch("homeassistant.components.goalzero.Yeti", return_value=mocked_yeti):
+        await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+    assert entry.state == ConfigEntryState.LOADED
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
+    assert entry.data == CONF_DATA
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.unique_id == "aa:bb:cc:dd:ee:ff"
+
+
 async def test_async_setup_entry_not_ready(hass: HomeAssistant) -> None:
     """Test that it throws ConfigEntryNotReady when exception occurs during setup."""
     entry = create_entry(hass)