Add detection helper

Author: CoMPaTech

airos/helpers.py

@@ -0,0 +1,71 @@
+"""Ubiquiti AirOS firmware helpers."""
+
+from typing import TypedDict
+
+import aiohttp
+
+from .airos6 import AirOS6
+from .airos8 import AirOS
+from .exceptions import AirOSKeyDataMissingError
+
+
+class DetectDeviceData(TypedDict):
+    """Container for device data."""
+
+    fw_major: int
+    mac: str
+    hostname: str
+
+
+async def async_get_firmware_data(
+    host: str,
+    username: str,
+    password: str,
+    session: aiohttp.ClientSession,
+    use_ssl: bool = True,
+) -> DetectDeviceData:
+    """Connect to a device and return the major firmware version."""
+    detect = AirOS(host, username, password, session, use_ssl)
+
+    await detect.login()
+    raw_status = await detect._request_json(  # noqa: SLF001
+        "GET",
+        detect._status_cgi_url,  # noqa: SLF001
+        authenticated=True,
+    )
+
+    fw_version = (raw_status.get("host") or {}).get("fwversion")
+    if not fw_version:
+        raise AirOSKeyDataMissingError("Missing host.fwversion in API response")
+
+    try:
+        fw_major = int(fw_version.lstrip("v").split(".", 1)[0])
+    except (ValueError, AttributeError) as exc:
+        raise AirOSKeyDataMissingError(
+            f"Invalid firmware version '{fw_version}'"
+        ) from exc
+
+    if fw_major == 6:
+        derived_data = AirOS6._derived_data_helper(  # noqa: SLF001
+            raw_status, AirOS6.derived_wireless_data
+        )
+    else:  # Assume AirOS 8 for all other versions
+        derived_data = AirOS._derived_data_helper(  # noqa: SLF001
+            raw_status, AirOS.derived_wireless_data
+        )
+
+    # Extract MAC address and hostname from the derived data
+    hostname = derived_data.get("host", {}).get("hostname")
+    mac = derived_data.get("derived", {}).get("mac")
+
+    if not hostname:
+        raise AirOSKeyDataMissingError("Missing hostname")
+
+    if not mac:
+        raise AirOSKeyDataMissingError("Missing MAC address")
+
+    return {
+        "fw_major": fw_major,
+        "mac": mac,
+        "hostname": hostname,
+    }

pyproject.toml

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

tests/test_helpers.py

@@ -0,0 +1,137 @@
+"""Test helpers for Ubiquiti airOS devices."""
+
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import aiohttp
+import pytest
+
+from airos.airos8 import AirOS
+from airos.exceptions import AirOSKeyDataMissingError
+from airos.helpers import DetectDeviceData, async_get_firmware_data
+
+# pylint: disable=redefined-outer-name
+
+
+@pytest.fixture
+def mock_session() -> MagicMock:
+    """Return a mock aiohttp ClientSession."""
+    return MagicMock(spec=aiohttp.ClientSession)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    (
+        "mock_response",
+        "expected_fw_major",
+        "expected_mac",
+        "expected_hostname",
+        "expected_exception",
+    ),
+    [
+        # Success case for AirOS 8
+        (
+            {
+                "host": {"fwversion": "v8.7.4", "hostname": "test-host-8"},
+                "interfaces": [
+                    {"hwaddr": "AA:BB:CC:DD:EE:FF", "ifname": "br0", "enabled": True}
+                ],
+            },
+            8,
+            "AA:BB:CC:DD:EE:FF",
+            "test-host-8",
+            None,
+        ),
+        # Success case for AirOS 6
+        (
+            {
+                "host": {"fwversion": "v6.3.16", "hostname": "test-host-6"},
+                "wireless": {"mode": "sta", "apmac": "11:22:33:44:55:66"},
+                "interfaces": [
+                    {"hwaddr": "11:22:33:44:55:66", "ifname": "br0", "enabled": True}
+                ],
+            },
+            6,
+            "11:22:33:44:55:66",
+            "test-host-6",
+            None,
+        ),
+        # Failure case: Missing host key
+        ({"wireless": {}}, 0, "", "", AirOSKeyDataMissingError),
+        # Failure case: Missing fwversion key
+        (
+            {"host": {"hostname": "test-host"}, "interfaces": []},
+            0,
+            "",
+            "",
+            AirOSKeyDataMissingError,
+        ),
+        # Failure case: Invalid fwversion value
+        (
+            {
+                "host": {"fwversion": "not-a-number", "hostname": "test-host"},
+                "interfaces": [],
+            },
+            0,
+            "",
+            "",
+            AirOSKeyDataMissingError,
+        ),
+        # Failure case: Missing hostname key
+        (
+            {"host": {"fwversion": "v8.7.4"}, "interfaces": []},
+            0,
+            "",
+            "",
+            AirOSKeyDataMissingError,
+        ),
+        # Failure case: Missing MAC address
+        (
+            {"host": {"fwversion": "v8.7.4", "hostname": "test-host"}},
+            0,
+            "",
+            "",
+            AirOSKeyDataMissingError,
+        ),
+    ],
+)
+async def test_firmware_detection(
+    mock_session: aiohttp.ClientSession,
+    mock_response: dict[str, Any],
+    expected_fw_major: int,
+    expected_mac: str,
+    expected_hostname: str,
+    expected_exception: Any,
+) -> None:
+    """Test helper firmware detection."""
+
+    mock_request_json = AsyncMock(
+        side_effect=[
+            {},  # First call for login()
+            mock_response,  # Second call for the status() endpoint
+        ]
+    )
+
+    with patch.object(AirOS, "_request_json", new=mock_request_json):
+        if expected_exception:
+            with pytest.raises(expected_exception):
+                await async_get_firmware_data(
+                    host="192.168.1.3",
+                    username="testuser",
+                    password="testpassword",
+                    session=mock_session,
+                    use_ssl=True,
+                )
+        else:
+            # Test the success case
+            device_data: DetectDeviceData = await async_get_firmware_data(
+                host="192.168.1.3",
+                username="testuser",
+                password="testpassword",
+                session=mock_session,
+                use_ssl=True,
+            )
+
+            assert device_data["fw_major"] == expected_fw_major
+            assert device_data["mac"] == expected_mac
+            assert device_data["hostname"] == expected_hostname