UniFi: Add UDM/P (UniFi OS) support (#33766)

* Fix get_controller and all tests:wq

* Bump dependency to v16

homeassistant/components/unifi/controller.py

@@ -327,6 +327,7 @@ async def get_controller(
 
     try:
         with async_timeout.timeout(10):
+            await controller.check_unifi_os()
             await controller.login()
         return controller
 

homeassistant/components/unifi/manifest.json

@@ -3,7 +3,7 @@
   "name": "Ubiquiti UniFi",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/unifi",
-  "requirements": ["aiounifi==15"],
+  "requirements": ["aiounifi==16"],
   "codeowners": ["@kane610"],
   "quality_scale": "platinum"
 }

requirements_all.txt

@@ -211,7 +211,7 @@ aiopylgtv==0.3.3
 aioswitcher==1.1.0
 
 # homeassistant.components.unifi
-aiounifi==15
+aiounifi==16
 
 # homeassistant.components.wwlln
 aiowwlln==2.0.2

requirements_test_all.txt

@@ -94,7 +94,7 @@ aiopylgtv==0.3.3
 aioswitcher==1.1.0
 
 # homeassistant.components.unifi
-aiounifi==15
+aiounifi==16
 
 # homeassistant.components.wwlln
 aiowwlln==2.0.2

tests/components/unifi/test_config_flow.py

@@ -52,6 +52,8 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
         CONF_VERIFY_SSL: False,
     }
 
+    aioclient_mock.get("https://1.2.3.4:1234", status=302)
+
     aioclient_mock.post(
         "https://1.2.3.4:1234/api/login",
         json={"data": "login successful", "meta": {"rc": "ok"}},
@@ -101,6 +103,8 @@ async def test_flow_works_multiple_sites(hass, aioclient_mock):
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "user"
 
+    aioclient_mock.get("https://1.2.3.4:1234", status=302)
+
     aioclient_mock.post(
         "https://1.2.3.4:1234/api/login",
         json={"data": "login successful", "meta": {"rc": "ok"}},
@@ -150,6 +154,8 @@ async def test_flow_fails_site_already_configured(hass, aioclient_mock):
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "user"
 
+    aioclient_mock.get("https://1.2.3.4:1234", status=302)
+
     aioclient_mock.post(
         "https://1.2.3.4:1234/api/login",
         json={"data": "login successful", "meta": {"rc": "ok"}},
@@ -188,6 +194,8 @@ async def test_flow_fails_user_credentials_faulty(hass, aioclient_mock):
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "user"
 
+    aioclient_mock.get("https://1.2.3.4:1234", status=302)
+
     with patch("aiounifi.Controller.login", side_effect=aiounifi.errors.Unauthorized):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -213,6 +221,8 @@ async def test_flow_fails_controller_unavailable(hass, aioclient_mock):
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "user"
 
+    aioclient_mock.get("https://1.2.3.4:1234", status=302)
+
     with patch("aiounifi.Controller.login", side_effect=aiounifi.errors.RequestError):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -238,6 +248,8 @@ async def test_flow_fails_unknown_problem(hass, aioclient_mock):
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "user"
 
+    aioclient_mock.get("https://1.2.3.4:1234", status=302)
+
     with patch("aiounifi.Controller.login", side_effect=Exception):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],

tests/components/unifi/test_controller.py

@@ -4,7 +4,7 @@ from copy import deepcopy
 from datetime import timedelta
 
 import aiounifi
-from asynctest import Mock, patch
+from asynctest import patch
 import pytest
 
 from homeassistant.components import unifi
@@ -68,11 +68,7 @@ async def setup_unifi_integration(
     controllers=None,
 ):
     """Create the UniFi controller."""
-    configuration = {}
+    assert await async_setup_component(hass, unifi.DOMAIN, {})
-    if controllers:
-        configuration = {unifi.DOMAIN: {unifi.CONF_CONTROLLERS: controllers}}
-
-    assert await async_setup_component(hass, unifi.DOMAIN, configuration)
 
     config_entry = MockConfigEntry(
         domain=unifi.DOMAIN,
@@ -108,20 +104,21 @@ async def setup_unifi_integration(
     async def mock_request(self, method, path, json=None):
         mock_requests.append({"method": method, "path": path, "json": json})
 
-        if path == "s/{site}/stat/sta" and mock_client_responses:
+        if path == "/stat/sta" and mock_client_responses:
             return mock_client_responses.popleft()
-        if path == "s/{site}/stat/device" and mock_device_responses:
+        if path == "/stat/device" and mock_device_responses:
             return mock_device_responses.popleft()
-        if path == "s/{site}/rest/user" and mock_client_all_responses:
+        if path == "/rest/user" and mock_client_all_responses:
             return mock_client_all_responses.popleft()
-        if path == "s/{site}/rest/wlanconf" and mock_wlans_responses:
+        if path == "/rest/wlanconf" and mock_wlans_responses:
             return mock_wlans_responses.popleft()
         return {}
 
-    # "aiounifi.Controller.start_websocket", return_value=True
+    with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
-    with patch("aiounifi.Controller.login", return_value=True), patch(
+        "aiounifi.Controller.login", return_value=True,
-        "aiounifi.Controller.sites", return_value=sites
+    ), patch("aiounifi.Controller.sites", return_value=sites), patch(
-    ), patch("aiounifi.Controller.request", new=mock_request), patch.object(
+        "aiounifi.Controller.request", new=mock_request
+    ), patch.object(
         aiounifi.websocket.WSClient, "start", return_value=True
     ):
         await hass.config_entries.async_setup(config_entry.entry_id)
@@ -244,7 +241,9 @@ async def test_wireless_client_event_calls_update_wireless_devices(hass):
 
 async def test_get_controller(hass):
     """Successful call."""
-    with patch("aiounifi.Controller.login", return_value=Mock()):
+    with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
+        "aiounifi.Controller.login", return_value=True
+    ):
         assert await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
 
 
@@ -252,13 +251,15 @@ async def test_get_controller_verify_ssl_false(hass):
     """Successful call with verify ssl set to false."""
     controller_data = dict(CONTROLLER_DATA)
     controller_data[CONF_VERIFY_SSL] = False
-    with patch("aiounifi.Controller.login", return_value=Mock()):
+    with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
+        "aiounifi.Controller.login", return_value=True
+    ):
         assert await unifi.controller.get_controller(hass, **controller_data)
 
 
 async def test_get_controller_login_failed(hass):
     """Check that get_controller can handle a failed login."""
-    with patch(
+    with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
         "aiounifi.Controller.login", side_effect=aiounifi.Unauthorized
     ), pytest.raises(unifi.errors.AuthenticationRequired):
         await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
@@ -266,7 +267,7 @@ async def test_get_controller_login_failed(hass):
 
 async def test_get_controller_controller_unavailable(hass):
     """Check that get_controller can handle controller being unavailable."""
-    with patch(
+    with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
         "aiounifi.Controller.login", side_effect=aiounifi.RequestError
     ), pytest.raises(unifi.errors.CannotConnect):
         await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
@@ -274,7 +275,7 @@ async def test_get_controller_controller_unavailable(hass):
 
 async def test_get_controller_unknown_error(hass):
     """Check that get_controller can handle unknown errors."""
-    with patch(
+    with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
         "aiounifi.Controller.login", side_effect=aiounifi.AiounifiException
     ), pytest.raises(unifi.errors.AuthenticationRequired):
         await unifi.controller.get_controller(hass, **CONTROLLER_DATA)

tests/components/unifi/test_switch.py

@@ -286,7 +286,7 @@ async def test_switches(hass):
     assert controller.mock_requests[4] == {
         "json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"},
         "method": "post",
-        "path": "s/{site}/cmd/stamgr/",
+        "path": "/cmd/stamgr",
     }
 
     await hass.services.async_call(
@@ -296,7 +296,7 @@ async def test_switches(hass):
     assert controller.mock_requests[5] == {
         "json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
         "method": "post",
-        "path": "s/{site}/cmd/stamgr/",
+        "path": "/cmd/stamgr",
     }
 
 
@@ -397,7 +397,7 @@ async def test_new_client_discovered_on_poe_control(hass):
             "port_overrides": [{"port_idx": 1, "portconf_id": "1a1", "poe_mode": "off"}]
         },
         "method": "put",
-        "path": "s/{site}/rest/device/mock-id",
+        "path": "/rest/device/mock-id",
     }
 
     await hass.services.async_call(
@@ -411,7 +411,7 @@ async def test_new_client_discovered_on_poe_control(hass):
             ]
         },
         "method": "put",
-        "path": "s/{site}/rest/device/mock-id",
+        "path": "/rest/device/mock-id",
     }
 
     switch_2 = hass.states.get("switch.poe_client_2")