Merge pull request #14 from CoMPaTech/testing

Add tests and station reconnect

Author: CoMPaTech

.gitignore

@@ -10,3 +10,4 @@ tests/__pycache__
 .vscode
 .coverage
 tmp
+todo

airos/airos8.py

@@ -4,7 +4,7 @@
 
 import json
 import logging
-from urllib.parse import urlparse
+from urllib.parse import quote, urlparse
 
 import aiohttp
 
@@ -47,6 +47,7 @@ def __init__(
 
         self._login_url = f"{self.base_url}/api/auth"  # AirOS 8
         self._status_cgi_url = f"{self.base_url}/status.cgi"  # AirOS 8
+        self._stakick_cgi_url = f"{self.base_url}/stakick.cgi"  # AirOS 8
         self.current_csrf_token = None
 
         self._use_json_for_login_post = False
@@ -147,19 +148,19 @@ async def login(self) -> bool:
                 # Re-check cookies in self.session.cookie_jar AFTER potential manual injection
                 airos_cookie_found = False
                 ok_cookie_found = False
-                if not self.session.cookie_jar:
+                if not self.session.cookie_jar:  # pragma: no cover
                     logger.exception(
                         "COOKIE JAR IS EMPTY after login POST. This is a major issue."
                     )
                     raise ConnectionSetupError from None
-                for cookie in self.session.cookie_jar:
+                for cookie in self.session.cookie_jar:  # pragma: no cover
                     if cookie.key.startswith("AIROS_"):
                         airos_cookie_found = True
                     if cookie.key == "ok":
                         ok_cookie_found = True
 
                 if not airos_cookie_found and not ok_cookie_found:
-                    raise ConnectionSetupError from None
+                    raise ConnectionSetupError from None  # pragma: no cover
 
                 response_text = await response.text()
 
@@ -176,7 +177,10 @@ async def login(self) -> bool:
                     log = f"Login failed with status {response.status}. Full Response: {response.text}"
                     logger.error(log)
                     raise ConnectionAuthenticationError from None
-        except aiohttp.ClientError as err:
+        except (
+            aiohttp.ClientError,
+            aiohttp.client_exceptions.ConnectionTimeoutError,
+        ) as err:
             logger.exception("Error during login")
             raise DeviceConnectionError from err
 
@@ -208,6 +212,49 @@ async def status(self) -> dict:
                 else:
                     log = f"Authenticated status.cgi failed: {response.status}. Response: {response_text}"
                     logger.error(log)
-        except aiohttp.ClientError as err:
+        except (
+            aiohttp.ClientError,
+            aiohttp.client_exceptions.ConnectionTimeoutError,
+        ) as err:
             logger.exception("Error during authenticated status.cgi call")
             raise DeviceConnectionError from err
+
+    async def stakick(self, mac_address: str = None) -> bool:
+        """Reconnect client station."""
+        if not self.connected:
+            logger.error("Not connected, login first")
+            raise DeviceConnectionError from None
+        if not mac_address:
+            logger.error("Device mac-address missing")
+            raise DataMissingError from None
+
+        # --- Step 2: Verify authenticated access by fetching status.cgi ---
+        kick_request_headers = {**self._common_headers}
+        if self.current_csrf_token:
+            kick_request_headers["X-CSRF-ID"] = self.current_csrf_token
+
+        kick_payload = {
+            "staif": "ath0",
+            "staid": quote(mac_address.upper(), safe=""),
+        }
+
+        kick_request_headers["Content-Type"] = (
+            "application/x-www-form-urlencoded; charset=UTF-8"
+        )
+        post_data = kick_payload
+
+        try:
+            async with self.session.post(
+                self._stakick_cgi_url,
+                headers=kick_request_headers,
+                data=post_data,
+            ) as response:
+                if response.status == 200:
+                    return True
+                return False
+        except (
+            aiohttp.ClientError,
+            aiohttp.client_exceptions.ConnectionTimeoutError,
+        ) as err:
+            logger.exception("Error during reconnect stakick.cgi call")
+            raise DeviceConnectionError from err

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "airos"
-version         = "0.0.8"
+version         = "0.0.9"
 license         = "MIT"
 description     = "Ubiquity airOS module(s) for Python 3."
 readme          = "README.md"

tests/test_stations.py

@@ -5,9 +5,11 @@
 import os
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import airos.exceptions
 import pytest
 
 import aiofiles
+import aiohttp
 
 
 async def _read_fixture(fixture: str = "ap-ptp"):
@@ -56,3 +58,88 @@ async def test_ap(airos_device, base_url, mode):
 
         # Verify the fixture returns the correct mode
         assert status.get("wireless", {}).get("mode") == mode
+
+
+@pytest.mark.asyncio
+async def test_reconnect(airos_device, base_url):
+    """Test reconnect client."""
+    # --- Prepare fake POST /api/stakick response ---
+    mock_stakick_response = MagicMock()
+    mock_stakick_response.__aenter__.return_value = mock_stakick_response
+    mock_stakick_response.status = 200
+
+    with (
+        patch.object(airos_device.session, "post", return_value=mock_stakick_response),
+        patch.object(airos_device, "connected", True),
+    ):
+        assert await airos_device.stakick("01:23:45:67:89:aB")
+
+
+@pytest.mark.asyncio
+async def test_ap_corners(airos_device, base_url, mode="ap-ptp"):
+    """Test device operation."""
+    cookie = SimpleCookie()
+    cookie["session_id"] = "test-cookie"
+    cookie["AIROS_TOKEN"] = "abc123"
+
+    # --- Prepare fake POST /api/auth response with cookies ---
+    mock_login_response = MagicMock()
+    mock_login_response.__aenter__.return_value = mock_login_response
+    mock_login_response.text = AsyncMock(return_value="{}")
+    mock_login_response.status = 200
+    mock_login_response.cookies = cookie
+    mock_login_response.headers = {"X-CSRF-ID": "test-csrf-token"}
+
+    with (
+        patch.object(airos_device.session, "post", return_value=mock_login_response),
+        patch.object(airos_device, "_use_json_for_login_post", return_value=True),
+    ):
+        assert await airos_device.login()
+
+    mock_login_response.cookies = {}
+    with (
+        patch.object(airos_device.session, "post", return_value=mock_login_response),
+    ):
+        try:
+            assert await airos_device.login()
+            assert False
+        except airos.exceptions.ConnectionSetupError:
+            assert True
+
+    mock_login_response.cookies = cookie
+    mock_login_response.headers = {}
+    with (
+        patch.object(airos_device.session, "post", return_value=mock_login_response),
+    ):
+        result = await airos_device.login()
+        assert result is None
+
+    mock_login_response.headers = {"X-CSRF-ID": "test-csrf-token"}
+    mock_login_response.text = AsyncMock(return_value="abc123")
+    with (
+        patch.object(airos_device.session, "post", return_value=mock_login_response),
+    ):
+        try:
+            assert await airos_device.login()
+            assert False
+        except airos.exceptions.DataMissingError:
+            assert True
+
+    mock_login_response.text = AsyncMock(return_value="{}")
+    mock_login_response.status = 400
+    with (
+        patch.object(airos_device.session, "post", return_value=mock_login_response),
+    ):
+        try:
+            assert await airos_device.login()
+            assert False
+        except airos.exceptions.ConnectionAuthenticationError:
+            assert True
+
+    mock_login_response.status = 200
+    with patch.object(airos_device.session, "post", side_effect=aiohttp.ClientError):
+        try:
+            assert await airos_device.login()
+            assert False
+        except airos.exceptions.DeviceConnectionError:
+            assert True