Migrate nut to use aionut (#114078)

This commit is contained in:
J. Nick Koston 2024-03-23 12:02:02 -10:00 committed by GitHub
parent d4f158d079
commit 4ac439ef88
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 123 additions and 107 deletions

View File

@ -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()

View File

@ -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}

View File

@ -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:

View File

@ -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."]
}

View File

@ -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

View File

@ -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

View File

@ -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(

View File

@ -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"},

View File

@ -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()

View File

@ -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)

View File

@ -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(