From 64277058b5ba6fb10029553422695964204f0ebb Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 17 Feb 2022 17:03:20 -0600
Subject: [PATCH] Ensure lutron caseta imports set the unique id (#66754)

---
 .../components/lutron_caseta/config_flow.py   | 28 ++++++++++++-------
 .../lutron_caseta/test_config_flow.py         |  6 ++++
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/lutron_caseta/config_flow.py b/homeassistant/components/lutron_caseta/config_flow.py
index b198d5ddbee..74819e25e8e 100644
--- a/homeassistant/components/lutron_caseta/config_flow.py
+++ b/homeassistant/components/lutron_caseta/config_flow.py
@@ -1,4 +1,6 @@
 """Config flow for Lutron Caseta."""
+from __future__ import annotations
+
 import asyncio
 import logging
 import os
@@ -17,6 +19,7 @@ from homeassistant.data_entry_flow import FlowResult
 
 from .const import (
     ABORT_REASON_CANNOT_CONNECT,
+    BRIDGE_DEVICE_ID,
     BRIDGE_TIMEOUT,
     CONF_CA_CERTS,
     CONF_CERTFILE,
@@ -101,7 +104,7 @@ class LutronCasetaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         if (
             not self.attempted_tls_validation
             and await self.hass.async_add_executor_job(self._tls_assets_exist)
-            and await self.async_validate_connectable_bridge_config()
+            and await self.async_get_lutron_id()
         ):
             self.tls_assets_validated = True
         self.attempted_tls_validation = True
@@ -177,7 +180,7 @@ class LutronCasetaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         self.data[CONF_CERTFILE] = import_info[CONF_CERTFILE]
         self.data[CONF_CA_CERTS] = import_info[CONF_CA_CERTS]
 
-        if not await self.async_validate_connectable_bridge_config():
+        if not (lutron_id := await self.async_get_lutron_id()):
             # Ultimately we won't have a dedicated step for import failure, but
             # in order to keep configuration.yaml-based configs transparently
             # working without requiring further actions from the user, we don't
@@ -189,6 +192,8 @@ class LutronCasetaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             # will require users to go through a confirmation flow for imports).
             return await self.async_step_import_failed()
 
+        await self.async_set_unique_id(lutron_id, raise_on_progress=False)
+        self._abort_if_unique_id_configured()
         return self.async_create_entry(title=ENTRY_DEFAULT_TITLE, data=self.data)
 
     async def async_step_import_failed(self, user_input=None):
@@ -204,10 +209,8 @@ class LutronCasetaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
         return self.async_abort(reason=ABORT_REASON_CANNOT_CONNECT)
 
-    async def async_validate_connectable_bridge_config(self):
+    async def async_get_lutron_id(self) -> str | None:
         """Check if we can connect to the bridge with the current config."""
-        bridge = None
-
         try:
             bridge = Smartbridge.create_tls(
                 hostname=self.data[CONF_HOST],
@@ -220,18 +223,23 @@ class LutronCasetaFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                 "Invalid certificate used to connect to bridge at %s",
                 self.data[CONF_HOST],
             )
-            return False
+            return None
 
-        connected_ok = False
         try:
             async with async_timeout.timeout(BRIDGE_TIMEOUT):
                 await bridge.connect()
-            connected_ok = bridge.is_connected()
         except asyncio.TimeoutError:
             _LOGGER.error(
                 "Timeout while trying to connect to bridge at %s",
                 self.data[CONF_HOST],
             )
+        else:
+            if not bridge.is_connected():
+                return None
+            devices = bridge.get_devices()
+            bridge_device = devices[BRIDGE_DEVICE_ID]
+            return hex(bridge_device["serial"])[2:].zfill(8)
+        finally:
+            await bridge.close()
 
-        await bridge.close()
-        return connected_ok
+        return None
diff --git a/tests/components/lutron_caseta/test_config_flow.py b/tests/components/lutron_caseta/test_config_flow.py
index 47956e27002..9dbedeacf5b 100644
--- a/tests/components/lutron_caseta/test_config_flow.py
+++ b/tests/components/lutron_caseta/test_config_flow.py
@@ -56,6 +56,10 @@ class MockBridge:
         """Return whether the mock bridge is connected."""
         return self.is_currently_connected
 
+    def get_devices(self):
+        """Return devices on the bridge."""
+        return {"1": {"serial": 1234}}
+
     async def close(self):
         """Close the mock bridge connection."""
         self.is_currently_connected = False
@@ -90,6 +94,8 @@ async def test_bridge_import_flow(hass):
     assert result["type"] == "create_entry"
     assert result["title"] == CasetaConfigFlow.ENTRY_DEFAULT_TITLE
     assert result["data"] == entry_mock_data
+    assert result["result"].unique_id == "000004d2"
+
     await hass.async_block_till_done()
     assert len(mock_setup_entry.mock_calls) == 1