From d0ed4b1ff26774a187ffb9b26fe0f58a8c6fcf51 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Wed, 19 Oct 2022 23:10:01 +0200 Subject: [PATCH] Replace constants with enums in UniFi (#80637) Replace constants with enums Fix bad imports --- homeassistant/components/unifi/controller.py | 2 +- .../components/unifi/device_tracker.py | 2 +- homeassistant/components/unifi/switch.py | 3 +- tests/components/unifi/conftest.py | 6 +- tests/components/unifi/test_controller.py | 13 ++-- tests/components/unifi/test_device_tracker.py | 59 +++++++++---------- tests/components/unifi/test_sensor.py | 10 ++-- tests/components/unifi/test_switch.py | 13 ++-- tests/components/unifi/test_update.py | 16 +++-- 9 files changed, 59 insertions(+), 65 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 6ad8e416445..1f3fd48140e 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -9,7 +9,7 @@ from typing import Any from aiohttp import CookieJar import aiounifi -from aiounifi.controller import ( +from aiounifi.interfaces.messages import ( DATA_CLIENT_REMOVED, DATA_DPI_GROUP, DATA_DPI_GROUP_REMOVED, diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index f184663291d..1c91ea4724b 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -3,7 +3,7 @@ from datetime import timedelta import logging -from aiounifi.interfaces.api_handlers import SOURCE_DATA, SOURCE_EVENT +from aiounifi.models.api import SOURCE_DATA, SOURCE_EVENT from aiounifi.models.event import EventKey from homeassistant.components.device_tracker import DOMAIN, SourceType diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index df935822c5a..68fe84b6b60 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -8,7 +8,8 @@ Support for controlling deep packet inspection (DPI) restriction groups. import asyncio from typing import Any -from aiounifi.interfaces.api_handlers import SOURCE_EVENT, ItemEvent +from aiounifi.interfaces.api_handlers import ItemEvent +from aiounifi.models.api import SOURCE_EVENT from aiounifi.models.client import ClientBlockRequest from aiounifi.models.device import ( DeviceSetOutletRelayRequest, diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index e2b77ac1ed1..40635823836 100644 --- a/tests/components/unifi/conftest.py +++ b/tests/components/unifi/conftest.py @@ -3,7 +3,7 @@ from __future__ import annotations from unittest.mock import patch -from aiounifi.websocket import SIGNAL_CONNECTION_STATE, SIGNAL_DATA +from aiounifi.websocket import WebsocketSignal import pytest from homeassistant.helpers import device_registry as dr @@ -20,10 +20,10 @@ def mock_unifi_websocket(): """Generate a websocket call.""" if data: mock.return_value.data = data - mock.call_args[1]["callback"](SIGNAL_DATA) + mock.call_args[1]["callback"](WebsocketSignal.DATA) elif state: mock.return_value.state = state - mock.call_args[1]["callback"](SIGNAL_CONNECTION_STATE) + mock.call_args[1]["callback"](WebsocketSignal.CONNECTION_STATE) else: raise NotImplementedError diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 5de99a3f434..0079e8e984c 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -7,7 +7,8 @@ from http import HTTPStatus from unittest.mock import Mock, patch import aiounifi -from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING +from aiounifi.models.event import EventKey +from aiounifi.websocket import WebsocketState import pytest from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN @@ -359,13 +360,13 @@ async def test_connection_state_signalling( # Controller is connected assert hass.states.get("device_tracker.client").state == "home" - mock_unifi_websocket(state=STATE_DISCONNECTED) + mock_unifi_websocket(state=WebsocketState.DISCONNECTED) await hass.async_block_till_done() # Controller is disconnected assert hass.states.get("device_tracker.client").state == "unavailable" - mock_unifi_websocket(state=STATE_RUNNING) + mock_unifi_websocket(state=WebsocketState.RUNNING) await hass.async_block_till_done() # Controller is once again connected @@ -403,7 +404,7 @@ async def test_wireless_client_event_calls_update_wireless_devices( { "datetime": "2020-01-20T19:37:04Z", "user": "00:00:00:00:00:01", - "key": aiounifi.events.WIRELESS_CLIENT_CONNECTED, + "key": EventKey.WIRELESS_CLIENT_CONNECTED.value, "msg": "User[11:22:33:44:55:66] has connected to WLAN", "time": 1579549024893, } @@ -423,7 +424,7 @@ async def test_reconnect_mechanism(hass, aioclient_mock, mock_unifi_websocket): f"https://{DEFAULT_HOST}:1234/api/login", status=HTTPStatus.BAD_GATEWAY ) - mock_unifi_websocket(state=STATE_DISCONNECTED) + mock_unifi_websocket(state=WebsocketState.DISCONNECTED) await hass.async_block_till_done() assert aioclient_mock.call_count == 0 @@ -459,7 +460,7 @@ async def test_reconnect_mechanism_exceptions( with patch("aiounifi.Controller.login", side_effect=exception), patch( "homeassistant.components.unifi.controller.UniFiController.reconnect" ) as mock_reconnect: - mock_unifi_websocket(state=STATE_DISCONNECTED) + mock_unifi_websocket(state=WebsocketState.DISCONNECTED) await hass.async_block_till_done() new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER) diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index e5d53a6d882..d004d96a73c 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -3,13 +3,8 @@ from datetime import timedelta from unittest.mock import patch -from aiounifi.controller import ( - MESSAGE_CLIENT, - MESSAGE_CLIENT_REMOVED, - MESSAGE_DEVICE, - MESSAGE_EVENT, -) -from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING +from aiounifi.models.message import MessageKey +from aiounifi.websocket import WebsocketState from homeassistant import config_entries from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN @@ -64,7 +59,7 @@ async def test_tracked_wireless_clients( client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -85,7 +80,7 @@ async def test_tracked_wireless_clients( mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -165,7 +160,7 @@ async def test_tracked_clients( client_1["last_seen"] += 1 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client_1], } ) @@ -215,7 +210,7 @@ async def test_tracked_wireless_clients_event_source( } mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_EVENT}, + "meta": {"message": MessageKey.EVENT.value}, "data": [event], } ) @@ -242,7 +237,7 @@ async def test_tracked_wireless_clients_event_source( } mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_EVENT}, + "meta": {"message": MessageKey.EVENT.value}, "data": [event], } ) @@ -265,7 +260,7 @@ async def test_tracked_wireless_clients_event_source( mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -292,7 +287,7 @@ async def test_tracked_wireless_clients_event_source( } mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_EVENT}, + "meta": {"message": MessageKey.EVENT.value}, "data": [event], } ) @@ -357,14 +352,14 @@ async def test_tracked_devices( device_1["next_interval"] = 20 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_DEVICE}, + "meta": {"message": MessageKey.DEVICE.value}, "data": [device_1], } ) device_2["next_interval"] = 50 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_DEVICE}, + "meta": {"message": MessageKey.DEVICE.value}, "data": [device_2], } ) @@ -388,7 +383,7 @@ async def test_tracked_devices( device_1["disabled"] = True mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_DEVICE}, + "meta": {"message": MessageKey.DEVICE.value}, "data": [device_1], } ) @@ -427,7 +422,7 @@ async def test_remove_clients( mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT_REMOVED}, + "meta": {"message": MessageKey.CLIENT_REMOVED.value}, "data": [client_1], } ) @@ -480,14 +475,14 @@ async def test_controller_state_change( assert hass.states.get("device_tracker.device").state == STATE_HOME # Controller unavailable - mock_unifi_websocket(state=STATE_DISCONNECTED) + mock_unifi_websocket(state=WebsocketState.DISCONNECTED) await hass.async_block_till_done() assert hass.states.get("device_tracker.client").state == STATE_UNAVAILABLE assert hass.states.get("device_tracker.device").state == STATE_UNAVAILABLE # Controller available - mock_unifi_websocket(state=STATE_RUNNING) + mock_unifi_websocket(state=WebsocketState.RUNNING) await hass.async_block_till_done() assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME @@ -730,7 +725,7 @@ async def test_option_ssid_filter( client["essid"] = "other_ssid" mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -738,7 +733,7 @@ async def test_option_ssid_filter( client_on_ssid2["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client_on_ssid2], } ) @@ -761,13 +756,13 @@ async def test_option_ssid_filter( client_on_ssid2["last_seen"] += 1 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client_on_ssid2], } ) @@ -788,7 +783,7 @@ async def test_option_ssid_filter( client_on_ssid2["last_seen"] += 1 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client_on_ssid2], } ) @@ -801,7 +796,7 @@ async def test_option_ssid_filter( client_on_ssid2["last_seen"] += 1 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client_on_ssid2], } ) @@ -850,7 +845,7 @@ async def test_wireless_client_go_wired_issue( client["is_wired"] = True mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -876,7 +871,7 @@ async def test_wireless_client_go_wired_issue( client["last_seen"] += 1 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -892,7 +887,7 @@ async def test_wireless_client_go_wired_issue( client["is_wired"] = False mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -936,7 +931,7 @@ async def test_option_ignore_wired_bug( client["is_wired"] = True mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -962,7 +957,7 @@ async def test_option_ignore_wired_bug( client["last_seen"] += 1 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) @@ -978,7 +973,7 @@ async def test_option_ignore_wired_bug( client["is_wired"] = False mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [client], } ) diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 398c6c6c3f5..0bf03b10029 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -3,7 +3,7 @@ from datetime import datetime from unittest.mock import patch -from aiounifi.controller import MESSAGE_CLIENT, MESSAGE_CLIENT_REMOVED +from aiounifi.models.message import MessageKey import pytest from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN @@ -89,7 +89,7 @@ async def test_bandwidth_sensors(hass, aioclient_mock, mock_unifi_websocket): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [wireless_client], } ) @@ -201,7 +201,7 @@ async def test_uptime_sensors( with patch("homeassistant.util.dt.now", return_value=now): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [uptime_client], } ) @@ -217,7 +217,7 @@ async def test_uptime_sensors( with patch("homeassistant.util.dt.now", return_value=now): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT}, + "meta": {"message": MessageKey.CLIENT.value}, "data": [uptime_client], } ) @@ -310,7 +310,7 @@ async def test_remove_sensors(hass, aioclient_mock, mock_unifi_websocket): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT_REMOVED}, + "meta": {"message": MessageKey.CLIENT_REMOVED.value}, "data": [wired_client], } ) diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 41494114c1e..e3ce7859e03 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -3,7 +3,6 @@ from copy import deepcopy from datetime import timedelta -from aiounifi.controller import MESSAGE_CLIENT_REMOVED, MESSAGE_DEVICE, MESSAGE_EVENT from aiounifi.models.message import MessageKey from aiounifi.websocket import WebsocketState @@ -745,7 +744,7 @@ async def test_remove_switches(hass, aioclient_mock, mock_unifi_websocket): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_CLIENT_REMOVED}, + "meta": {"message": MessageKey.CLIENT_REMOVED.value}, "data": [CLIENT_1, UNBLOCKED], } ) @@ -792,7 +791,7 @@ async def test_block_switches(hass, aioclient_mock, mock_unifi_websocket): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_EVENT}, + "meta": {"message": MessageKey.EVENT.value}, "data": [EVENT_BLOCKED_CLIENT_UNBLOCKED], } ) @@ -805,7 +804,7 @@ async def test_block_switches(hass, aioclient_mock, mock_unifi_websocket): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_EVENT}, + "meta": {"message": MessageKey.EVENT.value}, "data": [EVENT_BLOCKED_CLIENT_BLOCKED], } ) @@ -960,7 +959,7 @@ async def test_outlet_switches(hass, aioclient_mock, mock_unifi_websocket): mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_DEVICE}, + "meta": {"message": MessageKey.DEVICE.value}, "data": [outlet_up1], } ) @@ -1048,7 +1047,7 @@ async def test_new_client_discovered_on_block_control( mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_EVENT}, + "meta": {"message": MessageKey.EVENT.value}, "data": [EVENT_BLOCKED_CLIENT_CONNECTED], } ) @@ -1154,7 +1153,7 @@ async def test_new_client_discovered_on_poe_control( mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_EVENT}, + "meta": {"message": MessageKey.EVENT.value}, "data": [EVENT_CLIENT_2_CONNECTED], } ) diff --git a/tests/components/unifi/test_update.py b/tests/components/unifi/test_update.py index 677491319a1..9f212b5e065 100644 --- a/tests/components/unifi/test_update.py +++ b/tests/components/unifi/test_update.py @@ -1,8 +1,8 @@ """The tests for the UniFi Network update platform.""" from copy import deepcopy -from aiounifi.controller import MESSAGE_DEVICE -from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING +from aiounifi.models.message import MessageKey +from aiounifi.websocket import WebsocketState from yarl import URL from homeassistant.components.unifi.const import CONF_SITE_ID @@ -104,7 +104,7 @@ async def test_device_updates( device_1["state"] = 4 mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_DEVICE}, + "meta": {"message": MessageKey.DEVICE.value}, "data": [device_1], } ) @@ -124,7 +124,7 @@ async def test_device_updates( del device_1["upgrade_to_firmware"] mock_unifi_websocket( data={ - "meta": {"message": MESSAGE_DEVICE}, + "meta": {"message": MessageKey.DEVICE.value}, "data": [device_1], } ) @@ -188,9 +188,7 @@ async def test_install(hass, aioclient_mock): ) -async def test_controller_state_change( - hass, aioclient_mock, mock_unifi_websocket, mock_device_registry -): +async def test_controller_state_change(hass, aioclient_mock, mock_unifi_websocket): """Verify entities state reflect on controller becoming unavailable.""" await setup_unifi_integration( hass, @@ -202,13 +200,13 @@ async def test_controller_state_change( assert hass.states.get("update.device_1").state == STATE_ON # Controller unavailable - mock_unifi_websocket(state=STATE_DISCONNECTED) + mock_unifi_websocket(state=WebsocketState.DISCONNECTED) await hass.async_block_till_done() assert hass.states.get("update.device_1").state == STATE_UNAVAILABLE # Controller available - mock_unifi_websocket(state=STATE_RUNNING) + mock_unifi_websocket(state=WebsocketState.RUNNING) await hass.async_block_till_done() assert hass.states.get("update.device_1").state == STATE_ON