diff --git a/homeassistant/components/hyperion/__init__.py b/homeassistant/components/hyperion/__init__.py index 05494d14869..870d1e1bc3e 100644 --- a/homeassistant/components/hyperion/__init__.py +++ b/homeassistant/components/hyperion/__init__.py @@ -133,11 +133,13 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b ) and token is None ): + await hyperion_client.async_client_disconnect() await _create_reauth_flow(hass, config_entry) return False # Client login doesn't work? => Reauth. if not await hyperion_client.async_client_login(): + await hyperion_client.async_client_disconnect() await _create_reauth_flow(hass, config_entry) return False @@ -146,6 +148,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b not await hyperion_client.async_client_switch_instance() or not client.ServerInfoResponseOK(await hyperion_client.async_get_serverinfo()) ): + await hyperion_client.async_client_disconnect() raise ConfigEntryNotReady hyperion_client.set_callbacks( diff --git a/homeassistant/components/hyperion/light.py b/homeassistant/components/hyperion/light.py index e2989cb973b..e57d2916860 100644 --- a/homeassistant/components/hyperion/light.py +++ b/homeassistant/components/hyperion/light.py @@ -37,7 +37,11 @@ from homeassistant.helpers.typing import ( ) import homeassistant.util.color as color_util -from . import async_create_connect_hyperion_client, get_hyperion_unique_id +from . import ( + async_create_connect_hyperion_client, + create_hyperion_client, + get_hyperion_unique_id, +) from .const import ( CONF_ON_UNLOAD, CONF_PRIORITY, @@ -132,12 +136,12 @@ async def async_setup_platform( # First, connect to the server and get the server id (which will be unique_id on a config_entry # if there is one). - hyperion_client = await async_create_connect_hyperion_client(host, port) - if not hyperion_client: - raise PlatformNotReady - hyperion_id = await hyperion_client.async_sysinfo_id() - if not hyperion_id: - raise PlatformNotReady + async with create_hyperion_client(host, port) as hyperion_client: + if not hyperion_client: + raise PlatformNotReady + hyperion_id = await hyperion_client.async_sysinfo_id() + if not hyperion_id: + raise PlatformNotReady future_unique_id = get_hyperion_unique_id( hyperion_id, instance, TYPE_HYPERION_LIGHT diff --git a/tests/components/hyperion/test_light.py b/tests/components/hyperion/test_light.py index ea13aaab5ae..026cefcff69 100644 --- a/tests/components/hyperion/test_light.py +++ b/tests/components/hyperion/test_light.py @@ -105,6 +105,7 @@ async def test_setup_yaml_already_converted(hass: HomeAssistantType) -> None: add_test_config_entry(hass) client = create_mock_client() await _setup_entity_yaml(hass, client=client) + assert client.async_client_disconnect.called # Setup should be skipped for the YAML config as there is a pre-existing config # entry. @@ -127,6 +128,7 @@ async def test_setup_yaml_old_style_unique_id(hass: HomeAssistantType) -> None: client = create_mock_client() await _setup_entity_yaml(hass, client=client) + assert client.async_client_disconnect.called # The entity should have been created with the same entity_id. assert hass.states.get(TEST_YAML_ENTITY_ID) is not None @@ -166,6 +168,7 @@ async def test_setup_yaml_new_style_unique_id_wo_config( client = create_mock_client() await _setup_entity_yaml(hass, client=client) + assert client.async_client_disconnect.called # The entity should have been created with the same entity_id. assert hass.states.get(entity_id_to_preserve) is not None @@ -189,6 +192,7 @@ async def test_setup_yaml_no_registry_entity(hass: HomeAssistantType) -> None: # Add a pre-existing config entry. client = create_mock_client() await _setup_entity_yaml(hass, client=client) + assert client.async_client_disconnect.called # The entity should have been created with the same entity_id. assert hass.states.get(TEST_YAML_ENTITY_ID) is not None @@ -210,6 +214,7 @@ async def test_setup_yaml_not_ready(hass: HomeAssistantType) -> None: client = create_mock_client() client.async_client_connect = AsyncMock(return_value=False) await _setup_entity_yaml(hass, client=client) + assert client.async_client_disconnect.called assert hass.states.get(TEST_YAML_ENTITY_ID) is None @@ -236,6 +241,7 @@ async def test_setup_config_entry_not_ready_switch_instance_fail( client = create_mock_client() client.async_client_switch_instance = AsyncMock(return_value=False) await setup_test_config_entry(hass, hyperion_client=client) + assert client.async_client_disconnect.called assert hass.states.get(TEST_ENTITY_ID_1) is None @@ -252,6 +258,7 @@ async def test_setup_config_entry_not_ready_load_state_fail( ) await setup_test_config_entry(hass, hyperion_client=client) + assert client.async_client_disconnect.called assert hass.states.get(TEST_ENTITY_ID_1) is None @@ -802,6 +809,7 @@ async def test_setup_entry_no_token_reauth(hass: HomeAssistantType) -> None: "homeassistant.components.hyperion.client.HyperionClient", return_value=client ), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init: assert not await hass.config_entries.async_setup(config_entry.entry_id) + assert client.async_client_disconnect.called mock_flow_init.assert_called_once_with( DOMAIN, context={CONF_SOURCE: SOURCE_REAUTH}, @@ -825,6 +833,7 @@ async def test_setup_entry_bad_token_reauth(hass: HomeAssistantType) -> None: "homeassistant.components.hyperion.client.HyperionClient", return_value=client ), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init: assert not await hass.config_entries.async_setup(config_entry.entry_id) + assert client.async_client_disconnect.called mock_flow_init.assert_called_once_with( DOMAIN, context={CONF_SOURCE: SOURCE_REAUTH},