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 from __future__ import annotations
import asyncio
from dataclasses import dataclass from dataclasses import dataclass
from datetime import timedelta from datetime import timedelta
import logging 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.config_entries import ConfigEntry
from homeassistant.const import ( from homeassistant.const import (
@ -19,8 +18,9 @@ from homeassistant.const import (
CONF_RESOURCES, CONF_RESOURCES,
CONF_SCAN_INTERVAL, CONF_SCAN_INTERVAL,
CONF_USERNAME, CONF_USERNAME,
EVENT_HOMEASSISTANT_STOP,
) )
from homeassistant.core import HomeAssistant from homeassistant.core import Event, HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import device_registry as dr from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed 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) data = PyNUTData(host, port, alias, username, password)
entry.async_on_unload(data.async_shutdown)
async def async_update_data() -> dict[str, str]: async def async_update_data() -> dict[str, str]:
"""Fetch data from NUT.""" """Fetch data from NUT."""
async with asyncio.timeout(10): try:
await hass.async_add_executor_job(data.update) return await data.async_update()
if not data.status: except NUTError as err:
raise UpdateFailed("Error fetching UPS state") raise UpdateFailed(f"Error fetching UPS state: {err}") from err
return data.status
coordinator = DataUpdateCoordinator( coordinator = DataUpdateCoordinator(
hass, 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 # Fetch initial data so we have data when entities subscribe
await coordinator.async_config_entry_first_refresh() 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 status = coordinator.data
_LOGGER.debug("NUT Sensors Available: %s", status) _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: if username is not None and password is not None:
user_available_commands = { user_available_commands = {
device_supported_command 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 if device_supported_command in INTEGRATION_SUPPORTED_COMMANDS
} }
else: else:
@ -213,15 +220,14 @@ class PyNUTData:
alias: str | None, alias: str | None,
username: str | None, username: str | None,
password: str | None, password: str | None,
persistent: bool = True,
) -> None: ) -> None:
"""Initialize the data object.""" """Initialize the data object."""
self._host = host self._host = host
self._alias = alias self._alias = alias
# Establish client with persistent=False to open/close connection on self._client = AIONUTClient(self._host, port, username, password, 5, persistent)
# each update call. This is more reliable with async.
self._client = PyNUTClient(self._host, port, username, password, 5, False)
self.ups_list: dict[str, str] | None = None self.ups_list: dict[str, str] | None = None
self._status: dict[str, str] | None = None self._status: dict[str, str] | None = None
self._device_info: NUTDeviceInfo | None = None self._device_info: NUTDeviceInfo | None = None
@ -241,11 +247,11 @@ class PyNUTData:
"""Return the device info for the ups.""" """Return the device info for the ups."""
return self._device_info or NUTDeviceInfo() 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.""" """Get the ups alias from NUT."""
try: try:
ups_list: dict[str, str] = self._client.list_ups() ups_list = await self._client.list_ups()
except PyNUTError as err: except NUTError as err:
_LOGGER.error("Failure getting NUT ups alias, %s", err) _LOGGER.error("Failure getting NUT ups alias, %s", err)
return None return None
@ -268,42 +274,45 @@ class PyNUTData:
return device_info 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.""" """Get the ups status from NUT."""
if self._alias is None: 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: async def async_update(self) -> dict[str, str]:
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:
"""Fetch the latest status from NUT.""" """Fetch the latest status from NUT."""
self._status = self._get_status() self._status = await self._async_get_status()
if self._device_info is None: if self._device_info is None:
self._device_info = self._get_device_info() self._device_info = self._get_device_info()
return self._status
async def async_run_command( async def async_run_command(self, command_name: str) -> None:
self, hass: HomeAssistant, command_name: str | None
) -> None:
"""Invoke instant command in UPS.""" """Invoke instant command in UPS."""
if TYPE_CHECKING:
assert self._alias is not None
try: try:
await hass.async_add_executor_job( await self._client.run_command(self._alias, command_name)
self._client.run_command, self._alias, command_name except NUTError as err:
)
except PyNUTError as err:
raise HomeAssistantError( raise HomeAssistantError(
f"Error running command {command_name}, {err}" f"Error running command {command_name}, {err}"
) from 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.""" """Fetch the list of supported commands."""
if TYPE_CHECKING:
assert self._alias is not None
try: try:
return cast(dict[str, str], self._client.list_commands(self._alias)) return await self._client.list_commands(self._alias)
except PyNUTError as err: except NUTError as err:
_LOGGER.error("Error retrieving supported commands %s", err) _LOGGER.error("Error retrieving supported commands %s", err)
return None 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 import logging
from typing import Any from typing import Any
from aionut import NUTError
import voluptuous as vol import voluptuous as vol
from homeassistant.components import zeroconf 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) username = data.get(CONF_USERNAME)
password = data.get(CONF_PASSWORD) password = data.get(CONF_PASSWORD)
nut_data = PyNUTData(host, port, alias, username, password) nut_data = PyNUTData(host, port, alias, username, password, persistent=False)
await hass.async_add_executor_job(nut_data.update) try:
if not (status := nut_data.status): status = await nut_data.async_update()
raise CannotConnect 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} 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] device_id: str = config[CONF_DEVICE_ID]
entry_id = _get_entry_id_from_device_id(hass, device_id) entry_id = _get_entry_id_from_device_id(hass, device_id)
data: PyNUTData = hass.data[DOMAIN][entry_id][PYNUT_DATA] 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: def _get_device_action_name(command_name: str) -> str:

View file

@ -6,7 +6,7 @@
"documentation": "https://www.home-assistant.io/integrations/nut", "documentation": "https://www.home-assistant.io/integrations/nut",
"integration_type": "device", "integration_type": "device",
"iot_class": "local_polling", "iot_class": "local_polling",
"loggers": ["pynut2"], "loggers": ["aionut"],
"requirements": ["pynut2==2.1.2"], "requirements": ["aionut==4.0.0"],
"zeroconf": ["_nut._tcp.local."] "zeroconf": ["_nut._tcp.local."]
} }

View file

@ -314,6 +314,9 @@ aionanoleaf==0.2.1
# homeassistant.components.notion # homeassistant.components.notion
aionotion==2024.03.0 aionotion==2024.03.0
# homeassistant.components.nut
aionut==4.0.0
# homeassistant.components.oncue # homeassistant.components.oncue
aiooncue==0.3.7 aiooncue==0.3.7
@ -1987,9 +1990,6 @@ pynobo==1.6.0
# homeassistant.components.nuki # homeassistant.components.nuki
pynuki==1.6.3 pynuki==1.6.3
# homeassistant.components.nut
pynut2==2.1.2
# homeassistant.components.nws # homeassistant.components.nws
pynws==1.6.0 pynws==1.6.0

View file

@ -287,6 +287,9 @@ aionanoleaf==0.2.1
# homeassistant.components.notion # homeassistant.components.notion
aionotion==2024.03.0 aionotion==2024.03.0
# homeassistant.components.nut
aionut==4.0.0
# homeassistant.components.oncue # homeassistant.components.oncue
aiooncue==0.3.7 aiooncue==0.3.7
@ -1541,9 +1544,6 @@ pynobo==1.6.0
# homeassistant.components.nuki # homeassistant.components.nuki
pynuki==1.6.3 pynuki==1.6.3
# homeassistant.components.nut
pynut2==2.1.2
# homeassistant.components.nws # homeassistant.components.nws
pynws==1.6.0 pynws==1.6.0

View file

@ -3,7 +3,7 @@
from ipaddress import ip_address from ipaddress import ip_address
from unittest.mock import patch 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 import config_entries, data_entry_flow, setup
from homeassistant.components import zeroconf from homeassistant.components import zeroconf
@ -20,7 +20,7 @@ from homeassistant.const import (
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from .util import _get_mock_pynutclient from .util import _get_mock_nutclient
from tests.common import MockConfigEntry from tests.common import MockConfigEntry
@ -51,12 +51,12 @@ async def test_form_zeroconf(hass: HomeAssistant) -> None:
assert result["step_id"] == "user" assert result["step_id"] == "user"
assert result["errors"] == {} assert result["errors"] == {}
mock_pynut = _get_mock_pynutclient( mock_pynut = _get_mock_nutclient(
list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"] list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"]
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
), patch( ), patch(
"homeassistant.components.nut.async_setup_entry", "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["type"] == data_entry_flow.FlowResultType.FORM
assert result["errors"] == {} assert result["errors"] == {}
mock_pynut = _get_mock_pynutclient( mock_pynut = _get_mock_nutclient(
list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"] list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"]
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
), patch( ), patch(
"homeassistant.components.nut.async_setup_entry", "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["type"] == data_entry_flow.FlowResultType.FORM
assert result["errors"] == {} assert result["errors"] == {}
mock_pynut = _get_mock_pynutclient( mock_pynut = _get_mock_nutclient(
list_vars={"battery.voltage": "voltage"}, list_vars={"battery.voltage": "voltage"},
list_ups={"ups1": "UPS 1", "ups2": "UPS2"}, list_ups={"ups1": "UPS 1", "ups2": "UPS2"},
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
result2 = await hass.config_entries.flow.async_configure( 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 assert result2["type"] == data_entry_flow.FlowResultType.FORM
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
), patch( ), patch(
"homeassistant.components.nut.async_setup_entry", "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["type"] == data_entry_flow.FlowResultType.FORM
assert result["errors"] == {} assert result["errors"] == {}
mock_pynut = _get_mock_pynutclient( mock_pynut = _get_mock_nutclient(
list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"] list_vars={"battery.voltage": "voltage", "ups.status": "OL"}, list_ups=["ups1"]
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
), patch( ), patch(
"homeassistant.components.nut.async_setup_entry", "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} DOMAIN, context={"source": config_entries.SOURCE_USER}
) )
mock_pynut = _get_mock_pynutclient() mock_pynut = _get_mock_nutclient()
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
result2 = await hass.config_entries.flow.async_configure( 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"} assert result2["errors"] == {"base": "cannot_connect"}
with patch( with patch(
"homeassistant.components.nut.PyNUTClient.list_ups", "homeassistant.components.nut.AIONUTClient.list_ups",
side_effect=PyNUTError, side_effect=NUTError,
), patch( ), patch(
"homeassistant.components.nut.PyNUTClient.list_vars", "homeassistant.components.nut.AIONUTClient.list_vars",
side_effect=PyNUTError, side_effect=NUTError,
): ):
result2 = await hass.config_entries.flow.async_configure( result2 = await hass.config_entries.flow.async_configure(
result["flow_id"], result["flow_id"],
@ -278,11 +278,11 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None:
assert result2["errors"] == {"base": "cannot_connect"} assert result2["errors"] == {"base": "cannot_connect"}
with patch( with patch(
"homeassistant.components.nut.PyNUTClient.list_ups", "homeassistant.components.nut.AIONUTClient.list_ups",
return_value=["ups1"], return_value={"ups1"},
), patch( ), patch(
"homeassistant.components.nut.PyNUTClient.list_vars", "homeassistant.components.nut.AIONUTClient.list_vars",
side_effect=TypeError, side_effect=Exception,
): ):
result2 = await hass.config_entries.flow.async_configure( result2 = await hass.config_entries.flow.async_configure(
result["flow_id"], result["flow_id"],
@ -314,13 +314,13 @@ async def test_abort_if_already_setup(hass: HomeAssistant) -> None:
DOMAIN, context={"source": config_entries.SOURCE_USER} DOMAIN, context={"source": config_entries.SOURCE_USER}
) )
mock_pynut = _get_mock_pynutclient( mock_pynut = _get_mock_nutclient(
list_vars={"battery.voltage": "voltage"}, list_vars={"battery.voltage": "voltage"},
list_ups={"ups1": "UPS 1"}, list_ups={"ups1": "UPS 1"},
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
result2 = await hass.config_entries.flow.async_configure( 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} DOMAIN, context={"source": config_entries.SOURCE_USER}
) )
mock_pynut = _get_mock_pynutclient( mock_pynut = _get_mock_nutclient(
list_vars={"battery.voltage": "voltage"}, list_vars={"battery.voltage": "voltage"},
list_ups={"ups1": "UPS 1", "ups2": "UPS 2"}, list_ups={"ups1": "UPS 1", "ups2": "UPS 2"},
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
result2 = await hass.config_entries.flow.async_configure( 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 assert result2["type"] == data_entry_flow.FlowResultType.FORM
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
result3 = await hass.config_entries.flow.async_configure( result3 = await hass.config_entries.flow.async_configure(

View file

@ -1,8 +1,8 @@
"""The tests for Network UPS Tools (NUT) device actions.""" """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 import pytest
from pytest_unordered import unordered from pytest_unordered import unordered
@ -99,7 +99,7 @@ async def test_list_commands_exception(
) -> None: ) -> None:
"""Test there are no actions if list_commands raises exception.""" """Test there are no actions if list_commands raises exception."""
await async_init_integration( 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()) 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.enable": None,
"beeper.disable": None, "beeper.disable": None,
} }
run_command = MagicMock() run_command = AsyncMock()
await async_init_integration( await async_init_integration(
hass, hass,
list_ups={"someUps": "Some UPS"}, list_ups={"someUps": "Some UPS"},
@ -196,7 +196,7 @@ async def test_rund_command_exception(
list_commands_return_value = {"beeper.enable": None} list_commands_return_value = {"beeper.enable": None}
error_message = "Something wrong happened" 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( await async_init_integration(
hass, hass,
list_vars={"ups.status": "OL"}, list_vars={"ups.status": "OL"},

View file

@ -2,12 +2,14 @@
from unittest.mock import patch from unittest.mock import patch
from aionut import NUTError
from homeassistant.components.nut.const import DOMAIN from homeassistant.components.nut.const import DOMAIN
from homeassistant.config_entries import ConfigEntryState from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import CONF_HOST, CONF_PORT, STATE_UNAVAILABLE from homeassistant.const import CONF_HOST, CONF_PORT, STATE_UNAVAILABLE
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from .util import _get_mock_pynutclient from .util import _get_mock_nutclient
from tests.common import MockConfigEntry from tests.common import MockConfigEntry
@ -20,12 +22,12 @@ async def test_async_setup_entry(hass: HomeAssistant) -> None:
) )
entry.add_to_hass(hass) 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"} list_ups={"ups1": "UPS 1"}, list_vars={"ups.status": "OL"}
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
await hass.config_entries.async_setup(entry.entry_id) 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) entry.add_to_hass(hass)
with patch( with patch(
"homeassistant.components.nut.PyNUTClient.list_ups", "homeassistant.components.nut.AIONUTClient.list_ups",
return_value=["ups1"], return_value={"ups1"},
), patch( ), patch(
"homeassistant.components.nut.PyNUTClient.list_vars", "homeassistant.components.nut.AIONUTClient.list_vars",
side_effect=ConnectionResetError, side_effect=NUTError,
): ):
await hass.config_entries.async_setup(entry.entry_id) await hass.config_entries.async_setup(entry.entry_id)
await hass.async_block_till_done() await hass.async_block_till_done()

View file

@ -15,7 +15,7 @@ from homeassistant.const import (
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er 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 from tests.common import MockConfigEntry
@ -100,12 +100,12 @@ async def test_state_sensors(hass: HomeAssistant) -> None:
) )
entry.add_to_hass(hass) 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"} list_ups={"ups1": "UPS 1"}, list_vars={"ups.status": "OL"}
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
await hass.config_entries.async_setup(entry.entry_id) 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) 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"} list_ups={"ups1": "UPS 1"}, list_vars={"ups.status": "OQ"}
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
await hass.config_entries.async_setup(entry.entry_id) 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) 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"} list_ups={"ups1": "UPS 1"}, list_vars={"battery.charge": "10"}
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
await hass.config_entries.async_setup(config_entry.entry_id) await hass.config_entries.async_setup(config_entry.entry_id)

View file

@ -1,7 +1,7 @@
"""Tests for the nut integration.""" """Tests for the nut integration."""
import json import json
from unittest.mock import MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
from homeassistant.components.nut.const import DOMAIN from homeassistant.components.nut.const import DOMAIN
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME 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 from tests.common import MockConfigEntry, load_fixture
def _get_mock_pynutclient( def _get_mock_nutclient(
list_vars=None, list_vars=None,
list_ups=None, list_ups=None,
list_commands_return_value=None, list_commands_return_value=None,
list_commands_side_effect=None, list_commands_side_effect=None,
run_command=None, run_command=None,
): ):
pynutclient = MagicMock() nutclient = MagicMock()
type(pynutclient).list_ups = MagicMock(return_value=list_ups) type(nutclient).list_ups = AsyncMock(return_value=list_ups)
type(pynutclient).list_vars = MagicMock(return_value=list_vars) type(nutclient).list_vars = AsyncMock(return_value=list_vars)
if list_commands_return_value is None: if list_commands_return_value is None:
list_commands_return_value = {} 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 return_value=list_commands_return_value, side_effect=list_commands_side_effect
) )
if run_command is None: if run_command is None:
run_command = MagicMock() run_command = AsyncMock()
type(pynutclient).run_command = run_command type(nutclient).run_command = run_command
return pynutclient return nutclient
async def async_init_integration( async def async_init_integration(
@ -52,7 +52,7 @@ async def async_init_integration(
if list_vars is None: if list_vars is None:
list_vars = json.loads(load_fixture(ups_fixture)) list_vars = json.loads(load_fixture(ups_fixture))
mock_pynut = _get_mock_pynutclient( mock_pynut = _get_mock_nutclient(
list_ups=list_ups, list_ups=list_ups,
list_vars=list_vars, list_vars=list_vars,
list_commands_return_value=list_commands_return_value, list_commands_return_value=list_commands_return_value,
@ -61,7 +61,7 @@ async def async_init_integration(
) )
with patch( with patch(
"homeassistant.components.nut.PyNUTClient", "homeassistant.components.nut.AIONUTClient",
return_value=mock_pynut, return_value=mock_pynut,
): ):
entry = MockConfigEntry( entry = MockConfigEntry(