CRAI/Gem improvements

Author: CoMPaTech

README.md

@@ -142,19 +142,19 @@ if __name__ == "__main__":
 
 ##### Update
 
-Will return either ```{'update': False}``` or the full information regarding the available update:
+Will return either ```{"update": False}``` or the full information regarding the available update:
 
 ```json
-{'checksum': 'b1bea879a9f518f714ce638172e3a860', 'version': 'v8.7.19', 'security': '', 'date': '250811', 'url': 'https://dl.ubnt.com/firmwares/XC-fw/v8.7.19/WA.v8.7.19.48279.250811.0636.bin', 'update': True, 'changelog': 'https://dl.ubnt.com/firmwares/XC-fw/v8.7.19/changelog.txt'}
+{"checksum": "b1bea879a9f518f714ce638172e3a860", "version": "v8.7.19", "security": "", "date": "250811", "url": "https://dl.ubnt.com/firmwares/XC-fw/v8.7.19/WA.v8.7.19.48279.250811.0636.bin", "update": True, "changelog": "https://dl.ubnt.com/firmwares/XC-fw/v8.7.19/changelog.txt"}
 ```
 
 ##### Progress
 
-If no progress to report ```{'progress': -1}``` otherwise a positive value between 0 and 100.
+If no progress to report ```{"progress": -1}``` otherwise a positive value between 0 and 100.
 
 ##### Install
 
-Only positives outcome from user-experience, which should return:
+Only a positive outcome is expected from the user experience; the call should return:
 
 ```json
 {

airos/__init__.py

@@ -1 +1 @@
-"""Ubiquity airOS."""
+"""Ubiquiti airOS."""

airos/airos8.py

@@ -3,13 +3,15 @@
 from __future__ import annotations
 
 import asyncio
+from http.cookies import SimpleCookie
 import json
 import logging
-from typing import Any
+from typing import Any, NamedTuple
 from urllib.parse import urlparse
 
 import aiohttp
 from mashumaro.exceptions import InvalidFieldValue, MissingField
+from yarl import URL
 
 from .data import (
     AirOS8Data as AirOSData,
@@ -28,6 +30,16 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+class ApiResponse(NamedTuple):
+    """Define API call structure."""
+
+    status: int
+    headers: dict[str, Any]
+    cookies: SimpleCookie
+    url: URL
+    text: str
+
+
 class AirOS:
     """AirOS 8 connection class."""
 
@@ -159,7 +171,7 @@ def _get_authenticated_headers(
 
     async def _api_call(
         self, method: str, url: str, headers: dict[str, Any], **kwargs: Any
-    ) -> dict[str, Any]:
+    ) -> ApiResponse:
         """Make API call."""
         if url != self._login_url and not self.connected:
             _LOGGER.error("Not connected, login first")
@@ -170,8 +182,13 @@ async def _api_call(
                 method, url, headers=headers, **kwargs
             ) as response:
                 response_text = await response.text()
-                result = {"response": response, "response_text": response_text}
-                return result
+                return ApiResponse(
+                    status=response.status,
+                    headers=dict(response.headers),
+                    cookies=response.cookies,
+                    url=response.url,
+                    text=response_text,
+                )
         except (TimeoutError, aiohttp.ClientError) as err:
             _LOGGER.exception("Error during API call to %s: %s", url, err)
             raise AirOSDeviceConnectionError from err
@@ -183,9 +200,7 @@ async def _request_json(
         self, method: str, url: str, headers: dict[str, Any], **kwargs: Any
     ) -> dict[str, Any] | Any:
         """Return JSON from API call."""
-        result = await self._api_call(method, url, headers=headers, **kwargs)
-        response = result.get("response", {})
-        response_text = result.get("response_text", "")
+        response = await self._api_call(method, url, headers=headers, **kwargs)
 
         match response.status:
             case 200:
@@ -198,12 +213,12 @@ async def _request_json(
                     "API call to %s failed with status %d: %s",
                     url,
                     response.status,
-                    response_text,
+                    response.text,
                 )
                 raise AirOSDeviceConnectionError from None
 
         try:
-            return json.loads(await response.text())
+            return json.loads(response.text)
         except json.JSONDecodeError as err:
             _LOGGER.exception("JSON Decode Error in API response from %s", url)
             raise AirOSDataMissingError from err
@@ -222,26 +237,23 @@ async def login(self) -> bool:
         request_headers = self._get_authenticated_headers(ct_form=True)
         if self._use_json_for_login_post:
             request_headers = self._get_authenticated_headers(ct_json=True)
-            result = await self._api_call(
+            response = await self._api_call(
                 "POST", self._login_url, headers=request_headers, json=payload
             )
         else:
-            result = await self._api_call(
+            response = await self._api_call(
                 "POST", self._login_url, headers=request_headers, data=payload
             )
-        response = result.get("response", {})
-        response_text = result.get("response_text", "")
 
         if response.status == 403:
             _LOGGER.error("Authentication denied.")
             raise AirOSConnectionAuthenticationError from None
 
         for _, morsel in response.cookies.items():
             # If the AIROS_ cookie was parsed but isn't automatically added to the jar, add it manually
-            if (
-                morsel.key.startswith("AIROS_")
-                and morsel.key not in self.session.cookie_jar
-            ):
+            if morsel.key.startswith("AIROS_") and morsel.key not in [
+                cookie.key for cookie in self.session.cookie_jar
+            ]:
                 # `SimpleCookie`'s Morsel objects are designed to be compatible with cookie jars.
                 # We need to set the domain if it's missing, otherwise the cookie might not be sent.
                 # For IP addresses, the domain is typically blank.
@@ -299,7 +311,7 @@ async def login(self) -> bool:
             raise AirOSConnectionAuthenticationError from None
 
         try:
-            json.loads(response_text)
+            json.loads(response.text)
             self.connected = True
             return True
         except json.JSONDecodeError as err:
@@ -344,15 +356,13 @@ async def stakick(self, mac_address: str | None = None) -> bool:
         request_headers = self._get_authenticated_headers(ct_form=True)
         payload = {"staif": "ath0", "staid": mac_address.upper()}
 
-        result = await self._api_call(
+        response = await self._api_call(
             "POST", self._stakick_cgi_url, headers=request_headers, data=payload
         )
-        response = result.get("response", {})
         if response.status == 200:
             return True
 
-        response_text = result.get("response_text", "")
-        log = f"Unable to restart connection response status {response.status} with {response_text}"
+        log = f"Unable to restart connection response status {response.status} with {response.text}"
         _LOGGER.error(log)
         return False
 
@@ -365,15 +375,13 @@ async def provmode(self, active: bool = False) -> bool:
             action = "start"
 
         payload = {"action": action}
-        result = await self._api_call(
+        response = await self._api_call(
             "POST", self._provmode_url, headers=request_headers, data=payload
         )
-        response = result.get("response", {})
         if response.status == 200:
             return True
 
-        response_text = result.get("response_text", "")
-        log = f"Unable to change provisioning mode response status {response.status} with {response_text}"
+        log = f"Unable to change provisioning mode response status {response.status} with {response.text}"
         _LOGGER.error(log)
         return False
 
@@ -392,6 +400,10 @@ async def update_check(self, force: bool = False) -> dict[str, Any]:
         if force:
             payload = {"force": "yes"}
             request_headers = self._get_authenticated_headers(ct_form=True)
+            return await self._request_json(
+                "POST", self._update_check_url, headers=request_headers, data=payload
+            )
+
         return await self._request_json(
             "POST", self._update_check_url, headers=request_headers, json=payload
         )
@@ -402,7 +414,7 @@ async def progress(self) -> dict[str, Any]:
         payload: dict[str, Any] = {}
 
         return await self._request_json(
-            "POST", self._download_progress_url, headers=request_headers, json=payload
+            "POST", self._download_progress_url, headers=request_headers, data=payload
         )
 
     async def download(self) -> dict[str, Any]:

pyproject.toml

@@ -6,9 +6,9 @@ build-backend = "setuptools.build_meta"
 name            = "airos"
 version         = "0.4.0a0"
 license         = "MIT"
-description     = "Ubiquity airOS module(s) for Python 3."
+description     = "Ubiquiti airOS module(s) for Python 3."
 readme          = "README.md"
-keywords        = ["home", "automation", "ubiquity", "uisp", "airos", "module"]
+keywords        = ["home", "automation", "ubiquiti", "uisp", "airos", "module"]
 classifiers     = [
     "Development Status :: 5 - Production/Stable",
     "Intended Audience :: Developers",

tests/__init__.py

@@ -1 +1 @@
-"""Tests for the Ubiquity AirOS python module."""
+"""Tests for the Ubiquiti AirOS python module."""

tests/conftest.py

@@ -1,4 +1,4 @@
-"""Ubiquity AirOS test fixtures."""
+"""Ubiquiti AirOS test fixtures."""
 
 from _collections_abc import AsyncGenerator, Generator
 import asyncio

tests/test_stations.py

@@ -1,4 +1,4 @@
-"""Ubiquity AirOS tests."""
+"""Ubiquiti AirOS tests."""
 
 from http.cookies import SimpleCookie
 import json
@@ -12,6 +12,7 @@
 import pytest
 
 import aiofiles
+from yarl import URL
 
 
 async def _read_fixture(fixture: str = "loco5ac_ap-ptp") -> Any:
@@ -169,16 +170,19 @@ async def test_ap_object(
     mock_login_response.headers = {"X-CSRF-ID": "test-csrf-token"}
     # --- Prepare fake GET /api/status response ---
     fixture_data = await _read_fixture(fixture)
-    mock_status_payload = fixture_data
     mock_status_response = MagicMock()
     mock_status_response.__aenter__.return_value = mock_status_response
     mock_status_response.text = AsyncMock(return_value=json.dumps(fixture_data))
     mock_status_response.status = 200
-    mock_status_response.json = AsyncMock(return_value=mock_status_payload)
+    mock_status_response.cookies = SimpleCookie()
+    mock_status_response.headers = {}
+    mock_status_response.url = URL(base_url)
 
     with (
         patch.object(
-            airos_device.session, "request", return_value=mock_status_response
+            airos_device.session,
+            "request",
+            side_effect=[mock_login_response, mock_status_response],
         ),
     ):
         assert await airos_device.login()