diff --git a/homeassistant/components/esphome/manager.py b/homeassistant/components/esphome/manager.py index b897ffc9408..f197574c30a 100644 --- a/homeassistant/components/esphome/manager.py +++ b/homeassistant/components/esphome/manager.py @@ -383,6 +383,17 @@ class ESPHomeManager: self.voice_assistant_udp_server.stop() async def on_connect(self) -> None: + """Subscribe to states and list entities on successful API login.""" + try: + await self._on_connnect() + except APIConnectionError as err: + _LOGGER.warning( + "Error getting setting up connection for %s: %s", self.host, err + ) + # Re-connection logic will trigger after this + await self.cli.disconnect() + + async def _on_connnect(self) -> None: """Subscribe to states and list entities on successful API login.""" entry = self.entry unique_id = entry.unique_id @@ -393,16 +404,10 @@ class ESPHomeManager: cli = self.cli stored_device_name = entry.data.get(CONF_DEVICE_NAME) unique_id_is_mac_address = unique_id and ":" in unique_id - try: - results = await asyncio.gather( - cli.device_info(), - cli.list_entities_services(), - ) - except APIConnectionError as err: - _LOGGER.warning("Error getting device info for %s: %s", self.host, err) - # Re-connection logic will trigger after this - await cli.disconnect() - return + results = await asyncio.gather( + cli.device_info(), + cli.list_entities_services(), + ) device_info: EsphomeDeviceInfo = results[0] entity_infos_services: tuple[list[EntityInfo], list[UserService]] = results[1] @@ -487,18 +492,12 @@ class ESPHomeManager: ) ) - try: - setup_results = await asyncio.gather( - *setup_coros_with_disconnect_callbacks, - cli.subscribe_states(entry_data.async_update_state), - cli.subscribe_service_calls(self.async_on_service_call), - cli.subscribe_home_assistant_states(self.async_on_state_subscription), - ) - except APIConnectionError as err: - _LOGGER.warning("Error getting initial data for %s: %s", self.host, err) - # Re-connection logic will trigger after this - await cli.disconnect() - return + setup_results = await asyncio.gather( + *setup_coros_with_disconnect_callbacks, + cli.subscribe_states(entry_data.async_update_state), + cli.subscribe_service_calls(self.async_on_service_call), + cli.subscribe_home_assistant_states(self.async_on_state_subscription), + ) for result_idx in range(len(setup_coros_with_disconnect_callbacks)): cancel_callback = setup_results[result_idx] diff --git a/tests/components/esphome/conftest.py b/tests/components/esphome/conftest.py index 0ac940018d7..8c46fac08d4 100644 --- a/tests/components/esphome/conftest.py +++ b/tests/components/esphome/conftest.py @@ -90,6 +90,10 @@ class BaseMockReconnectLogic(ReconnectLogic): self._cancel_connect("forced disconnect from test") self._is_stopped = True + async def stop(self) -> None: + """Stop the reconnect logic.""" + self.stop_callback() + @pytest.fixture def mock_device_info() -> DeviceInfo: diff --git a/tests/components/esphome/test_manager.py b/tests/components/esphome/test_manager.py index 96a8a341308..474b1a757fb 100644 --- a/tests/components/esphome/test_manager.py +++ b/tests/components/esphome/test_manager.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, call from aioesphomeapi import ( APIClient, + APIConnectionError, DeviceInfo, EntityInfo, EntityState, @@ -510,8 +511,11 @@ async def test_connection_aborted_wrong_device( "with mac address `11:22:33:44:55:ab`" in caplog.text ) + assert "Error getting setting up connection for" not in caplog.text + assert len(mock_client.disconnect.mock_calls) == 1 + mock_client.disconnect.reset_mock() caplog.clear() - # Make sure discovery triggers a reconnect to the correct device + # Make sure discovery triggers a reconnect service_info = dhcp.DhcpServiceInfo( ip="192.168.43.184", hostname="test", @@ -533,6 +537,35 @@ async def test_connection_aborted_wrong_device( assert "Unexpected device found at" not in caplog.text +async def test_failure_during_connect( + hass: HomeAssistant, + mock_client: APIClient, + mock_zeroconf: None, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test we disconnect when there is a failure during connection setup.""" + entry = MockConfigEntry( + domain=DOMAIN, + data={ + CONF_HOST: "192.168.43.183", + CONF_PORT: 6053, + CONF_PASSWORD: "", + CONF_DEVICE_NAME: "test", + }, + unique_id="11:22:33:44:55:aa", + ) + entry.add_to_hass(hass) + + mock_client.device_info = AsyncMock(side_effect=APIConnectionError("fail")) + + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + + assert "Error getting setting up connection for" in caplog.text + # Ensure we disconnect so that the reconnect logic is triggered + assert len(mock_client.disconnect.mock_calls) == 1 + + async def test_state_subscription( mock_client: APIClient, hass: HomeAssistant,