From 8aec26d1051848a65a9de15bd7c39a23b97a35d1 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 18 Jan 2022 12:38:16 -1000
Subject: [PATCH] Cleanup unifiprotect discovery name (#64389)

---
 .../components/unifiprotect/config_flow.py    |  4 +-
 .../components/unifiprotect/services.py       |  7 +-
 .../components/unifiprotect/utils.py          | 14 ++++
 .../unifiprotect/test_config_flow.py          | 64 +++++++++++++++++++
 4 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/unifiprotect/config_flow.py b/homeassistant/components/unifiprotect/config_flow.py
index 3055f236696..7f06539c2fa 100644
--- a/homeassistant/components/unifiprotect/config_flow.py
+++ b/homeassistant/components/unifiprotect/config_flow.py
@@ -35,7 +35,7 @@ from .const import (
     OUTDATED_LOG_MESSAGE,
 )
 from .discovery import async_start_discovery
-from .services import _async_unifi_mac_from_hass
+from .utils import _async_short_mac, _async_unifi_mac_from_hass
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -97,7 +97,7 @@ class ProtectFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         placeholders = {
             "name": discovery_info["hostname"]
             or discovery_info["platform"]
-            or f"NVR {discovery_info['mac']}",
+            or f"NVR {_async_short_mac(discovery_info['mac'])}",
             "ip_address": discovery_info["ip_address"],
         }
         self.context["title_placeholders"] = placeholders
diff --git a/homeassistant/components/unifiprotect/services.py b/homeassistant/components/unifiprotect/services.py
index 3e58961b752..a149f516d2f 100644
--- a/homeassistant/components/unifiprotect/services.py
+++ b/homeassistant/components/unifiprotect/services.py
@@ -19,6 +19,7 @@ from homeassistant.helpers.service import async_extract_referenced_entity_ids
 
 from .const import ATTR_MESSAGE, DOMAIN
 from .data import ProtectData
+from .utils import _async_unifi_mac_from_hass
 
 SERVICE_ADD_DOORBELL_TEXT = "add_doorbell_text"
 SERVICE_REMOVE_DOORBELL_TEXT = "remove_doorbell_text"
@@ -48,12 +49,6 @@ def _async_all_ufp_instances(hass: HomeAssistant) -> list[ProtectApiClient]:
     ]
 
 
-@callback
-def _async_unifi_mac_from_hass(mac: str) -> str:
-    # MAC addresses in UFP are always caps
-    return mac.replace(":", "").upper()
-
-
 @callback
 def _async_get_macs_for_device(device_entry: dr.DeviceEntry) -> list[str]:
     return [
diff --git a/homeassistant/components/unifiprotect/utils.py b/homeassistant/components/unifiprotect/utils.py
index 14a64798ca2..45645d6f06b 100644
--- a/homeassistant/components/unifiprotect/utils.py
+++ b/homeassistant/components/unifiprotect/utils.py
@@ -4,6 +4,8 @@ from __future__ import annotations
 from enum import Enum
 from typing import Any
 
+from homeassistant.core import callback
+
 
 def get_nested_attr(obj: Any, attr: str) -> Any:
     """Fetch a nested attribute."""
@@ -19,3 +21,15 @@ def get_nested_attr(obj: Any, attr: str) -> Any:
         value = value.value
 
     return value
+
+
+@callback
+def _async_unifi_mac_from_hass(mac: str) -> str:
+    # MAC addresses in UFP are always caps
+    return mac.replace(":", "").upper()
+
+
+@callback
+def _async_short_mac(mac: str) -> str:
+    """Get the short mac address from the full mac."""
+    return _async_unifi_mac_from_hass(mac)[-6:]
diff --git a/tests/components/unifiprotect/test_config_flow.py b/tests/components/unifiprotect/test_config_flow.py
index 8194b05e060..b162c7a9540 100644
--- a/tests/components/unifiprotect/test_config_flow.py
+++ b/tests/components/unifiprotect/test_config_flow.py
@@ -51,6 +51,12 @@ UNIFI_DISCOVERY_DICT = {
     "hostname": DEVICE_HOSTNAME,
     "platform": DEVICE_HOSTNAME,
 }
+UNIFI_DISCOVERY_DICT_PARTIAL = {
+    "ip_address": DEVICE_IP_ADDRESS,
+    "mac": DEVICE_MAC_ADDRESS,
+    "hostname": None,
+    "platform": None,
+}
 
 
 async def test_form(hass: HomeAssistant, mock_nvr: NVR) -> None:
@@ -301,6 +307,64 @@ async def test_discovered_by_unifi_discovery(
 
     assert result["type"] == RESULT_TYPE_FORM
     assert result["step_id"] == "discovery_confirm"
+    flows = hass.config_entries.flow.async_progress_by_handler(DOMAIN)
+    assert flows[0]["context"]["title_placeholders"] == {
+        "ip_address": DEVICE_IP_ADDRESS,
+        "name": DEVICE_HOSTNAME,
+    }
+
+    assert not result["errors"]
+
+    with patch(
+        "homeassistant.components.unifiprotect.config_flow.ProtectApiClient.get_nvr",
+        return_value=mock_nvr,
+    ), patch(
+        "homeassistant.components.unifiprotect.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "username": "test-username",
+                "password": "test-password",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result2["type"] == RESULT_TYPE_CREATE_ENTRY
+    assert result2["title"] == "UnifiProtect"
+    assert result2["data"] == {
+        "host": DEVICE_IP_ADDRESS,
+        "username": "test-username",
+        "password": "test-password",
+        "id": "UnifiProtect",
+        "port": 443,
+        "verify_ssl": False,
+    }
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_discovered_by_unifi_discovery_partial(
+    hass: HomeAssistant, mock_nvr: NVR
+) -> None:
+    """Test a discovery from unifi-discovery partial."""
+
+    with _patch_discovery():
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DISCOVERY},
+            data=UNIFI_DISCOVERY_DICT_PARTIAL,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == "discovery_confirm"
+    flows = hass.config_entries.flow.async_progress_by_handler(DOMAIN)
+    assert flows[0]["context"]["title_placeholders"] == {
+        "ip_address": DEVICE_IP_ADDRESS,
+        "name": "NVR DDEEFF",
+    }
+
     assert not result["errors"]
 
     with patch(