Connect concurrently to discovered Zerproc lights (#44376)

* Connect concurrently to discovered Zerproc lights

* Add return type to connect_light

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/zerproc/light.py

@@ -1,4 +1,5 @@
 """Zerproc light platform."""
+import asyncio
 from datetime import timedelta
 import logging
 from typing import Callable, List, Optional
@@ -29,6 +30,16 @@ SUPPORT_ZERPROC = SUPPORT_BRIGHTNESS | SUPPORT_COLOR
 DISCOVERY_INTERVAL = timedelta(seconds=60)
 
 
+async def connect_light(light: pyzerproc.Light) -> Optional[pyzerproc.Light]:
+    """Return the given light if it connects successfully."""
+    try:
+        await light.connect()
+    except pyzerproc.ZerprocException:
+        _LOGGER.debug("Unable to connect to '%s'", light.address, exc_info=True)
+        return None
+    return light
+
+
 async def discover_entities(hass: HomeAssistant) -> List[Entity]:
     """Attempt to discover new lights."""
     lights = await pyzerproc.discover()
@@ -39,12 +50,10 @@ async def discover_entities(hass: HomeAssistant) -> List[Entity]:
     ]
 
     entities = []
-    for light in new_lights:
-        try:
-            await light.connect()
-        except pyzerproc.ZerprocException:
-            _LOGGER.debug("Unable to connect to '%s'", light.address, exc_info=True)
-            continue
+    connected_lights = filter(
+        None, await asyncio.gather(*(connect_light(light) for light in new_lights))
+    )
+    for light in connected_lights:
         # Double-check the light hasn't been added in the meantime
         if light.address not in hass.data[DOMAIN]["addresses"]:
             hass.data[DOMAIN]["addresses"].add(light.address)

tests/components/zerproc/test_light.py

@@ -141,22 +141,28 @@ async def test_connect_exception(hass, mock_entry):
 
     mock_entry.add_to_hass(hass)
 
-    mock_light = MagicMock(spec=pyzerproc.Light)
-    mock_light.address = "AA:BB:CC:DD:EE:FF"
-    mock_light.name = "LEDBlue-CCDDEEFF"
-    mock_light.is_connected.return_value = False
+    mock_light_1 = MagicMock(spec=pyzerproc.Light)
+    mock_light_1.address = "AA:BB:CC:DD:EE:FF"
+    mock_light_1.name = "LEDBlue-CCDDEEFF"
+    mock_light_1.is_connected.return_value = False
+
+    mock_light_2 = MagicMock(spec=pyzerproc.Light)
+    mock_light_2.address = "11:22:33:44:55:66"
+    mock_light_2.name = "LEDBlue-33445566"
+    mock_light_2.is_connected.return_value = False
 
     with patch(
         "homeassistant.components.zerproc.light.pyzerproc.discover",
-        return_value=[mock_light],
+        return_value=[mock_light_1, mock_light_2],
     ), patch.object(
-        mock_light, "connect", side_effect=pyzerproc.ZerprocException("TEST")
+        mock_light_1, "connect", side_effect=pyzerproc.ZerprocException("TEST")
     ):
         await hass.config_entries.async_setup(mock_entry.entry_id)
         await hass.async_block_till_done()
 
-    # The exception should be captured and no entities should be added
-    assert len(hass.data[DOMAIN]["addresses"]) == 0
+    # The exception connecting to light 1 should be captured, but light 2
+    # should still be added
+    assert len(hass.data[DOMAIN]["addresses"]) == 1
 
 
 async def test_remove_entry(hass, mock_light, mock_entry):