Add retries for tplink discovery (#52015)

* Add retries for tplink discovery

* Black Format tplink common.py

* Exit tplink discovery early if all devices found

* Fix typo in tplink retry log msg

* Code style cleanup for tplink retries

* Update homeassistant/components/tplink/common.py

Co-authored-by: Teemu R. <tpr@iki.fi>

* Fix linting errors for tplink retries

Co-authored-by: Teemu R. <tpr@iki.fi>

homeassistant/components/tplink/__init__.py

@@ -7,6 +7,7 @@ from homeassistant import config_entries
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.typing import ConfigType
 
@@ -76,6 +77,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up TPLink from a config entry."""
     config_data = hass.data[DOMAIN].get(ATTR_CONFIG)
 
+    device_registry = dr.async_get(hass)
+    tplink_devices = dr.async_entries_for_config_entry(device_registry, entry.entry_id)
+    device_count = len(tplink_devices)
+
     # These will contain the initialized devices
     lights = hass.data[DOMAIN][CONF_LIGHT] = []
     switches = hass.data[DOMAIN][CONF_SWITCH] = []
@@ -90,7 +95,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # Add discovered devices
     if config_data is None or config_data[CONF_DISCOVERY]:
-        discovered_devices = await async_discover_devices(hass, static_devices)
+        discovered_devices = await async_discover_devices(
+            hass, static_devices, device_count
+        )
 
         lights.extend(discovered_devices.lights)
         switches.extend(discovered_devices.switches)

homeassistant/components/tplink/common.py

@@ -26,6 +26,7 @@ CONF_DISCOVERY = "discovery"
 CONF_LIGHT = "light"
 CONF_STRIP = "strip"
 CONF_SWITCH = "switch"
+MAX_DISCOVERY_RETRIES = 4
 
 
 class SmartDevices:
@@ -67,12 +68,9 @@ async def async_get_discoverable_devices(hass: HomeAssistant) -> dict[str, Smart
 
 
 async def async_discover_devices(
-    hass: HomeAssistant, existing_devices: SmartDevices
+    hass: HomeAssistant, existing_devices: SmartDevices, target_device_count: int
 ) -> SmartDevices:
     """Get devices through discovery."""
-    _LOGGER.debug("Discovering devices")
-    devices = await async_get_discoverable_devices(hass)
-    _LOGGER.info("Discovered %s TP-Link smart home device(s)", len(devices))
 
     lights = []
     switches = []
@@ -100,6 +98,33 @@ async def async_discover_devices(
             else:
                 _LOGGER.error("Unknown smart device type: %s", type(dev))
 
+    devices = {}
+    for attempt in range(1, MAX_DISCOVERY_RETRIES + 1):
+        _LOGGER.debug(
+            "Discovering tplink devices, attempt %s of %s",
+            attempt,
+            MAX_DISCOVERY_RETRIES,
+        )
+        discovered_devices = await async_get_discoverable_devices(hass)
+        _LOGGER.info(
+            "Discovered %s TP-Link of expected %s smart home device(s)",
+            len(discovered_devices),
+            target_device_count,
+        )
+        for device_ip in discovered_devices:
+            devices[device_ip] = discovered_devices[device_ip]
+
+        if len(discovered_devices) >= target_device_count:
+            _LOGGER.info(
+                "Discovered at least as many devices on the network as exist in our device registry, no need to retry"
+            )
+            break
+
+    _LOGGER.info(
+        "Found %s unique TP-Link smart home device(s) after %s discovery attempts",
+        len(devices),
+        attempt,
+    )
     await hass.async_add_executor_job(process_devices)
 
     return SmartDevices(lights, switches)