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>
This commit is contained in:
Emily Mills 2020-12-19 13:31:45 -06:00 committed by GitHub
parent 60ecc8282c
commit 4f088dd77a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 29 additions and 14 deletions

View File

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

View File

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