Fix WiZ Light config flow timeout by properly closing UDP connections

homeassistant/components/wiz/config_flow.py

@@ -81,6 +81,8 @@ async def _async_connect_discovered_or_abort(self) -> None:
                 exc_info=True,
             )
             raise AbortFlow("cannot_connect") from ex
+        finally:
+            await bulb.async_close()
         self._name = name_from_bulb_type_and_mac(bulbtype, device.mac_address)
 
     async def async_step_discovery_confirm(
@@ -118,6 +120,8 @@ async def async_step_pick_device(
                 bulbtype = await bulb.get_bulbtype()
             except WIZ_CONNECT_EXCEPTIONS:
                 return self.async_abort(reason="cannot_connect")
+            finally:
+                await bulb.async_close()
 
             return self.async_create_entry(
                 title=name_from_bulb_type_and_mac(bulbtype, device.mac_address),
@@ -182,6 +186,8 @@ async def async_step_user(
                         title=name,
                         data=user_input,
                     )
+                finally:
+                    await bulb.async_close()
 
         return self.async_show_form(
             step_id="user",

tests/components/wiz/test_config_flow.py

@@ -47,14 +47,16 @@
 
 async def test_form(hass: HomeAssistant) -> None:
     """Test we get the form."""
+    bulb = _mocked_wizlight(None, None, FAKE_DIMMABLE_BULB)
+
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
     assert result["type"] is FlowResultType.FORM
     assert result["errors"] == {}
     # Patch functions
     with (
-        _patch_wizlight(),
+        _patch_wizlight(device=bulb),
         patch(
             "homeassistant.components.wiz.async_setup_entry",
             return_value=True,
@@ -76,6 +78,7 @@ async def test_form(hass: HomeAssistant) -> None:
     }
     assert len(mock_setup.mock_calls) == 1
     assert len(mock_setup_entry.mock_calls) == 1
+    bulb.async_close.assert_awaited_once()
 
 
 async def test_user_flow_enters_dns_name(hass: HomeAssistant) -> None:
@@ -137,17 +140,18 @@ async def test_user_form_exceptions(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
 
-    with patch(
-        "homeassistant.components.wiz.wizlight.getBulbConfig",
-        side_effect=side_effect,
-    ):
+    bulb = _mocked_wizlight(None, None, FAKE_DIMMABLE_BULB)
+    bulb.get_bulbtype = AsyncMock(side_effect=side_effect)
+
+    with _patch_wizlight(device=bulb):
         result2 = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             TEST_CONNECTION,
         )
 
     assert result2["type"] is FlowResultType.FORM
     assert result2["errors"] == {"base": error_base}
+    bulb.async_close.assert_awaited_once()
 
 
 async def test_form_updates_unique_id(hass: HomeAssistant) -> None:
@@ -185,17 +189,18 @@ async def test_discovered_by_dhcp_connection_fails(
     hass: HomeAssistant, source, data
 ) -> None:
     """Test we abort on connection failure."""
-    with patch(
-        "homeassistant.components.wiz.wizlight.getBulbConfig",
-        side_effect=WizLightTimeOutError,
-    ):
+    bulb = _mocked_wizlight(None, None, FAKE_DIMMABLE_BULB)
+    bulb.get_bulbtype = AsyncMock(side_effect=WizLightTimeOutError)
+
+    with _patch_wizlight(device=bulb):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": source}, data=data
         )
         await hass.async_block_till_done()
 
     assert result["type"] is FlowResultType.ABORT
     assert result["reason"] == "cannot_connect"
+    bulb.async_close.assert_awaited_once()
 
 
 @pytest.mark.parametrize(
@@ -263,21 +268,22 @@ async def test_discovered_by_dhcp_or_integration_discovery(
     hass: HomeAssistant, source, data, bulb_type, extended_white_range, name
 ) -> None:
     """Test we can configure when discovered from dhcp or discovery."""
-    with _patch_wizlight(
-        device=None, extended_white_range=extended_white_range, bulb_type=bulb_type
-    ):
+    bulb = _mocked_wizlight(None, extended_white_range, bulb_type)
+
+    with _patch_wizlight(device=bulb):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": source}, data=data
         )
         await hass.async_block_till_done()
 
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "discovery_confirm"
+    bulb.async_close.assert_awaited_once()
+
+    bulb.async_close.reset_mock()
 
     with (
-        _patch_wizlight(
-            device=None, extended_white_range=extended_white_range, bulb_type=bulb_type
-        ),
+        _patch_wizlight(device=bulb),
         patch(
             "homeassistant.components.wiz.async_setup_entry",
             return_value=True,
@@ -299,6 +305,7 @@ async def test_discovered_by_dhcp_or_integration_discovery(
     }
     assert len(mock_setup.mock_calls) == 1
     assert len(mock_setup_entry.mock_calls) == 1
+    bulb.async_close.assert_awaited_once()
 
 
 @pytest.mark.parametrize(
@@ -393,8 +400,10 @@ async def test_setup_via_discovery(hass: HomeAssistant) -> None:
     assert result2["step_id"] == "pick_device"
     assert not result2["errors"]
 
+    bulb = _mocked_wizlight(None, None, FAKE_DIMMABLE_BULB)
+
     with (
-        _patch_wizlight(),
+        _patch_wizlight(device=bulb),
         patch(
             "homeassistant.components.wiz.async_setup", return_value=True
         ) as mock_setup,
@@ -415,6 +424,7 @@ async def test_setup_via_discovery(hass: HomeAssistant) -> None:
     }
     assert len(mock_setup.mock_calls) == 1
     assert len(mock_setup_entry.mock_calls) == 1
+    bulb.async_close.assert_awaited_once()
 
     # ignore configured devices
     result = await hass.config_entries.flow.async_init(