diff --git a/homeassistant/components/nut/__init__.py b/homeassistant/components/nut/__init__.py index d5a2453f62f..c9067bbb254 100644 --- a/homeassistant/components/nut/__init__.py +++ b/homeassistant/components/nut/__init__.py @@ -2,13 +2,12 @@ from __future__ import annotations -import asyncio from dataclasses import dataclass from datetime import timedelta import logging -from typing import cast +from typing import TYPE_CHECKING -from pynut2.nut2 import PyNUTClient, PyNUTError +from aionut import AIONUTClient, NUTError from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( @@ -19,8 +18,9 @@ from homeassistant.const import ( CONF_RESOURCES, CONF_SCAN_INTERVAL, CONF_USERNAME, + EVENT_HOMEASSISTANT_STOP, ) -from homeassistant.core import HomeAssistant +from homeassistant.core import Event, HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import device_registry as dr from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed @@ -64,13 +64,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: data = PyNUTData(host, port, alias, username, password) + entry.async_on_unload(data.async_shutdown) + async def async_update_data() -> dict[str, str]: """Fetch data from NUT.""" - async with asyncio.timeout(10): - await hass.async_add_executor_job(data.update) - if not data.status: - raise UpdateFailed("Error fetching UPS state") - return data.status + try: + return await data.async_update() + except NUTError as err: + raise UpdateFailed(f"Error fetching UPS state: {err}") from err coordinator = DataUpdateCoordinator( hass, @@ -83,6 +84,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # Fetch initial data so we have data when entities subscribe await coordinator.async_config_entry_first_refresh() + + # Note that async_listen_once is not used here because the listener + # could be removed after the event is fired. + entry.async_on_unload( + hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, data.async_shutdown) + ) status = coordinator.data _LOGGER.debug("NUT Sensors Available: %s", status) @@ -95,7 +102,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if username is not None and password is not None: user_available_commands = { device_supported_command - for device_supported_command in data.list_commands() or {} + for device_supported_command in await data.async_list_commands() or {} if device_supported_command in INTEGRATION_SUPPORTED_COMMANDS } else: @@ -213,15 +220,14 @@ class PyNUTData: alias: str | None, username: str | None, password: str | None, + persistent: bool = True, ) -> None: """Initialize the data object.""" self._host = host self._alias = alias - # Establish client with persistent=False to open/close connection on - # each update call. This is more reliable with async. - self._client = PyNUTClient(self._host, port, username, password, 5, False) + self._client = AIONUTClient(self._host, port, username, password, 5, persistent) self.ups_list: dict[str, str] | None = None self._status: dict[str, str] | None = None self._device_info: NUTDeviceInfo | None = None @@ -241,11 +247,11 @@ class PyNUTData: """Return the device info for the ups.""" return self._device_info or NUTDeviceInfo() - def _get_alias(self) -> str | None: + async def _async_get_alias(self) -> str | None: """Get the ups alias from NUT.""" try: - ups_list: dict[str, str] = self._client.list_ups() - except PyNUTError as err: + ups_list = await self._client.list_ups() + except NUTError as err: _LOGGER.error("Failure getting NUT ups alias, %s", err) return None @@ -268,42 +274,45 @@ class PyNUTData: return device_info - def _get_status(self) -> dict[str, str] | None: + async def _async_get_status(self) -> dict[str, str]: """Get the ups status from NUT.""" if self._alias is None: - self._alias = self._get_alias() + self._alias = await self._async_get_alias() + if TYPE_CHECKING: + assert self._alias is not None + return await self._client.list_vars(self._alias) - try: - status: dict[str, str] = self._client.list_vars(self._alias) - except (PyNUTError, ConnectionResetError) as err: - _LOGGER.debug("Error getting NUT vars for host %s: %s", self._host, err) - return None - - return status - - def update(self) -> None: + async def async_update(self) -> dict[str, str]: """Fetch the latest status from NUT.""" - self._status = self._get_status() + self._status = await self._async_get_status() if self._device_info is None: self._device_info = self._get_device_info() + return self._status - async def async_run_command( - self, hass: HomeAssistant, command_name: str | None - ) -> None: + async def async_run_command(self, command_name: str) -> None: """Invoke instant command in UPS.""" + if TYPE_CHECKING: + assert self._alias is not None + try: - await hass.async_add_executor_job( - self._client.run_command, self._alias, command_name - ) - except PyNUTError as err: + await self._client.run_command(self._alias, command_name) + except NUTError as err: raise HomeAssistantError( f"Error running command {command_name}, {err}" ) from err - def list_commands(self) -> dict[str, str] | None: + async def async_list_commands(self) -> set[str] | None: """Fetch the list of supported commands.""" + if TYPE_CHECKING: + assert self._alias is not None + try: - return cast(dict[str, str], self._client.list_commands(self._alias)) - except PyNUTError as err: + return await self._client.list_commands(self._alias) + except NUTError as err: _LOGGER.error("Error retrieving supported commands %s", err) return None + + @callback + def async_shutdown(self, _: Event | None = None) -> None: + """Shutdown the client connection.""" + self._client.shutdown() diff --git a/homeassistant/components/nut/config_flow.py b/homeassistant/components/nut/config_flow.py index 5967aef4dac..3f3de8a126c 100644 --- a/homeassistant/components/nut/config_flow.py +++ b/homeassistant/components/nut/config_flow.py @@ -6,6 +6,7 @@ from collections.abc import Mapping import logging from typing import Any +from aionut import NUTError import voluptuous as vol from homeassistant.components import zeroconf @@ -67,10 +68,14 @@ async def validate_input(hass: HomeAssistant, data: dict[str, Any]) -> dict[str, username = data.get(CONF_USERNAME) password = data.get(CONF_PASSWORD) - nut_data = PyNUTData(host, port, alias, username, password) - await hass.async_add_executor_job(nut_data.update) - if not (status := nut_data.status): - raise CannotConnect + nut_data = PyNUTData(host, port, alias, username, password, persistent=False) + try: + status = await nut_data.async_update() + except NUTError as err: + raise CannotConnect(str(err)) from err + + if not alias and not nut_data.ups_list: + raise CannotConnect("No UPSes found on the NUT server") return {"ups_list": nut_data.ups_list, "available_resources": status} diff --git a/homeassistant/components/nut/device_action.py b/homeassistant/components/nut/device_action.py index bf16b3664f1..0ec58e651b2 100644 --- a/homeassistant/components/nut/device_action.py +++ b/homeassistant/components/nut/device_action.py @@ -58,7 +58,7 @@ async def async_call_action_from_config( device_id: str = config[CONF_DEVICE_ID] entry_id = _get_entry_id_from_device_id(hass, device_id) data: PyNUTData = hass.data[DOMAIN][entry_id][PYNUT_DATA] - await data.async_run_command(hass, command_name) + await data.async_run_command(command_name) def _get_device_action_name(command_name: str) -> str: diff --git a/homeassistant/components/nut/manifest.json b/homeassistant/components/nut/manifest.json index 0303dd70ec1..6cd073a7476 100644 --- a/homeassistant/components/nut/manifest.json +++ b/homeassistant/components/nut/manifest.json @@ -6,7 +6,7 @@ "documentation": "https://www.home-assistant.io/integrations/nut", "integration_type": "device", "iot_class": "local_polling", - "loggers": ["pynut2"], - "requirements": ["pynut2==2.1.2"], + "loggers": ["aionut"], + "requirements": ["aionut==4.0.0"], "zeroconf": ["_nut._tcp.local."] } diff --git a/requirements_all.txt b/requirements_all.txt index 27cdf5a353f..7ca3f82dae1 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -314,6 +314,9 @@ aionanoleaf==0.2.1 # homeassistant.components.notion aionotion==2024.03.0 +# homeassistant.components.nut +aionut==4.0.0 + # homeassistant.components.oncue aiooncue==0.3.7 @@ -1987,9 +1990,6 @@ pynobo==1.6.0 # homeassistant.components.nuki pynuki==1.6.3 -# homeassistant.components.nut -pynut2==2.1.2 - # homeassistant.components.nws pynws==1.6.0 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 8ccb0f1dd24..eee4500495d 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -287,6 +287,9 @@ aionanoleaf==0.2.1 # homeassistant.components.notion aionotion==2024.03.0 +# homeassistant.components.nut +aionut==4.0.0 + # homeassistant.components.oncue aiooncue==0.3.7 @@ -1541,9 +1544,6 @@ pynobo==1.6.0 # homeassistant.components.nuki pynuki==1.6.3 -# homeassistant.components.nut -pynut2==2.1.2 - # homeassistant.components.nws pynws==1.6.0 diff --git a/tests/components/nut/test_config_flow.py b/tests/components/nut/test_config_flow.py index 6dca92d1294..b6a9590f457 100644 --- a/tests/components/nut/test_config_flow.py +++ b/tests/components/nut/test_config_flow.py @@ -3,7 +3,7 @@ from ipaddress import ip_address from unittest.mock import patch -from pynut2.nut2 import PyNUTError +from aionut import NUTError from homeassistant import config_entries, data_entry_flow, setup from homeassistant.components import zeroconf @@ -20,7 +20,7 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant -from .util import _get_mock_pynutclient +from .util import _get_mock_nutclient from tests.common import MockConfigEntry @@ -51,12 +51,12 @@ async def test_form_zeroconf(hass: HomeAssistant) -> None: assert result["step_id"] == "user" assert result["errors"] == {} - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"] ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ), patch( "homeassistant.components.nut.async_setup_entry", @@ -89,12 +89,12 @@ async def test_form_user_one_ups(hass: HomeAssistant) -> None: assert result["type"] == data_entry_flow.FlowResultType.FORM assert result["errors"] == {} - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"] ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ), patch( "homeassistant.components.nut.async_setup_entry", @@ -138,13 +138,13 @@ async def test_form_user_multiple_ups(hass: HomeAssistant) -> None: assert result["type"] == data_entry_flow.FlowResultType.FORM assert result["errors"] == {} - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_vars={"battery.voltage": "voltage"}, list_ups={"ups1": "UPS 1", "ups2": "UPS2"}, ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): result2 = await hass.config_entries.flow.async_configure( @@ -161,7 +161,7 @@ async def test_form_user_multiple_ups(hass: HomeAssistant) -> None: assert result2["type"] == data_entry_flow.FlowResultType.FORM with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ), patch( "homeassistant.components.nut.async_setup_entry", @@ -199,12 +199,12 @@ async def test_form_user_one_ups_with_ignored_entry(hass: HomeAssistant) -> None assert result["type"] == data_entry_flow.FlowResultType.FORM assert result["errors"] == {} - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"] ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ), patch( "homeassistant.components.nut.async_setup_entry", @@ -238,10 +238,10 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None: DOMAIN, context={"source": config_entries.SOURCE_USER} ) - mock_pynut = _get_mock_pynutclient() + mock_pynut = _get_mock_nutclient() with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): result2 = await hass.config_entries.flow.async_configure( @@ -258,11 +258,11 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None: assert result2["errors"] == {"base": "cannot_connect"} with patch( - "homeassistant.components.nut.PyNUTClient.list_ups", - side_effect=PyNUTError, + "homeassistant.components.nut.AIONUTClient.list_ups", + side_effect=NUTError, ), patch( - "homeassistant.components.nut.PyNUTClient.list_vars", - side_effect=PyNUTError, + "homeassistant.components.nut.AIONUTClient.list_vars", + side_effect=NUTError, ): result2 = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -278,11 +278,11 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None: assert result2["errors"] == {"base": "cannot_connect"} with patch( - "homeassistant.components.nut.PyNUTClient.list_ups", - return_value=["ups1"], + "homeassistant.components.nut.AIONUTClient.list_ups", + return_value={"ups1"}, ), patch( - "homeassistant.components.nut.PyNUTClient.list_vars", - side_effect=TypeError, + "homeassistant.components.nut.AIONUTClient.list_vars", + side_effect=Exception, ): result2 = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -314,13 +314,13 @@ async def test_abort_if_already_setup(hass: HomeAssistant) -> None: DOMAIN, context={"source": config_entries.SOURCE_USER} ) - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_vars={"battery.voltage": "voltage"}, list_ups={"ups1": "UPS 1"}, ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): result2 = await hass.config_entries.flow.async_configure( @@ -352,13 +352,13 @@ async def test_abort_if_already_setup_alias(hass: HomeAssistant) -> None: DOMAIN, context={"source": config_entries.SOURCE_USER} ) - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_vars={"battery.voltage": "voltage"}, list_ups={"ups1": "UPS 1", "ups2": "UPS 2"}, ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): result2 = await hass.config_entries.flow.async_configure( @@ -373,7 +373,7 @@ async def test_abort_if_already_setup_alias(hass: HomeAssistant) -> None: assert result2["type"] == data_entry_flow.FlowResultType.FORM with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): result3 = await hass.config_entries.flow.async_configure( diff --git a/tests/components/nut/test_device_action.py b/tests/components/nut/test_device_action.py index 8938de54457..8113b19e313 100644 --- a/tests/components/nut/test_device_action.py +++ b/tests/components/nut/test_device_action.py @@ -1,8 +1,8 @@ """The tests for Network UPS Tools (NUT) device actions.""" -from unittest.mock import MagicMock +from unittest.mock import AsyncMock -from pynut2.nut2 import PyNUTError +from aionut import NUTError import pytest from pytest_unordered import unordered @@ -99,7 +99,7 @@ async def test_list_commands_exception( ) -> None: """Test there are no actions if list_commands raises exception.""" await async_init_integration( - hass, list_vars={"ups.status": "OL"}, list_commands_side_effect=PyNUTError + hass, list_vars={"ups.status": "OL"}, list_commands_side_effect=NUTError ) device_entry = next(device for device in device_registry.devices.values()) @@ -137,7 +137,7 @@ async def test_action(hass: HomeAssistant, device_registry: dr.DeviceRegistry) - "beeper.enable": None, "beeper.disable": None, } - run_command = MagicMock() + run_command = AsyncMock() await async_init_integration( hass, list_ups={"someUps": "Some UPS"}, @@ -196,7 +196,7 @@ async def test_rund_command_exception( list_commands_return_value = {"beeper.enable": None} error_message = "Something wrong happened" - run_command = MagicMock(side_effect=PyNUTError(error_message)) + run_command = AsyncMock(side_effect=NUTError(error_message)) await async_init_integration( hass, list_vars={"ups.status": "OL"}, diff --git a/tests/components/nut/test_init.py b/tests/components/nut/test_init.py index 2cee229897d..d15e9d4b12a 100644 --- a/tests/components/nut/test_init.py +++ b/tests/components/nut/test_init.py @@ -2,12 +2,14 @@ from unittest.mock import patch +from aionut import NUTError + from homeassistant.components.nut.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.const import CONF_HOST, CONF_PORT, STATE_UNAVAILABLE from homeassistant.core import HomeAssistant -from .util import _get_mock_pynutclient +from .util import _get_mock_nutclient from tests.common import MockConfigEntry @@ -20,12 +22,12 @@ async def test_async_setup_entry(hass: HomeAssistant) -> None: ) entry.add_to_hass(hass) - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_ups={"ups1": "UPS 1"}, list_vars={"ups.status": "OL"} ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): await hass.config_entries.async_setup(entry.entry_id) @@ -55,11 +57,11 @@ async def test_config_not_ready(hass: HomeAssistant) -> None: entry.add_to_hass(hass) with patch( - "homeassistant.components.nut.PyNUTClient.list_ups", - return_value=["ups1"], + "homeassistant.components.nut.AIONUTClient.list_ups", + return_value={"ups1"}, ), patch( - "homeassistant.components.nut.PyNUTClient.list_vars", - side_effect=ConnectionResetError, + "homeassistant.components.nut.AIONUTClient.list_vars", + side_effect=NUTError, ): await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() diff --git a/tests/components/nut/test_sensor.py b/tests/components/nut/test_sensor.py index be435fb6e6c..c4a8159b8cc 100644 --- a/tests/components/nut/test_sensor.py +++ b/tests/components/nut/test_sensor.py @@ -15,7 +15,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er -from .util import _get_mock_pynutclient, async_init_integration +from .util import _get_mock_nutclient, async_init_integration from tests.common import MockConfigEntry @@ -100,12 +100,12 @@ async def test_state_sensors(hass: HomeAssistant) -> None: ) entry.add_to_hass(hass) - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_ups={"ups1": "UPS 1"}, list_vars={"ups.status": "OL"} ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): await hass.config_entries.async_setup(entry.entry_id) @@ -125,12 +125,12 @@ async def test_unknown_state_sensors(hass: HomeAssistant) -> None: ) entry.add_to_hass(hass) - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_ups={"ups1": "UPS 1"}, list_vars={"ups.status": "OQ"} ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): await hass.config_entries.async_setup(entry.entry_id) @@ -155,12 +155,12 @@ async def test_stale_options(hass: HomeAssistant) -> None: ) config_entry.add_to_hass(hass) - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_ups={"ups1": "UPS 1"}, list_vars={"battery.charge": "10"} ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): await hass.config_entries.async_setup(config_entry.entry_id) diff --git a/tests/components/nut/util.py b/tests/components/nut/util.py index a0fadf47f19..3bc48764816 100644 --- a/tests/components/nut/util.py +++ b/tests/components/nut/util.py @@ -1,7 +1,7 @@ """Tests for the nut integration.""" import json -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from homeassistant.components.nut.const import DOMAIN from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME @@ -10,25 +10,25 @@ from homeassistant.core import HomeAssistant from tests.common import MockConfigEntry, load_fixture -def _get_mock_pynutclient( +def _get_mock_nutclient( list_vars=None, list_ups=None, list_commands_return_value=None, list_commands_side_effect=None, run_command=None, ): - pynutclient = MagicMock() - type(pynutclient).list_ups = MagicMock(return_value=list_ups) - type(pynutclient).list_vars = MagicMock(return_value=list_vars) + nutclient = MagicMock() + type(nutclient).list_ups = AsyncMock(return_value=list_ups) + type(nutclient).list_vars = AsyncMock(return_value=list_vars) if list_commands_return_value is None: list_commands_return_value = {} - type(pynutclient).list_commands = MagicMock( + type(nutclient).list_commands = AsyncMock( return_value=list_commands_return_value, side_effect=list_commands_side_effect ) if run_command is None: - run_command = MagicMock() - type(pynutclient).run_command = run_command - return pynutclient + run_command = AsyncMock() + type(nutclient).run_command = run_command + return nutclient async def async_init_integration( @@ -52,7 +52,7 @@ async def async_init_integration( if list_vars is None: list_vars = json.loads(load_fixture(ups_fixture)) - mock_pynut = _get_mock_pynutclient( + mock_pynut = _get_mock_nutclient( list_ups=list_ups, list_vars=list_vars, list_commands_return_value=list_commands_return_value, @@ -61,7 +61,7 @@ async def async_init_integration( ) with patch( - "homeassistant.components.nut.PyNUTClient", + "homeassistant.components.nut.AIONUTClient", return_value=mock_pynut, ): entry = MockConfigEntry(