From 7f6fadacf459f388b10b32b95b0beb80166a1e70 Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:37:35 +0200 Subject: [PATCH 1/8] Adjust min cov --- .github/workflows/verify.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/verify.yml b/.github/workflows/verify.yml index f8e785a..ade99eb 100644 --- a/.github/workflows/verify.yml +++ b/.github/workflows/verify.yml @@ -151,7 +151,7 @@ jobs: run: | . venv/bin/activate coverage combine coverage*/.coverage* - coverage report --fail-under=85 + coverage report --fail-under=80 coverage xml - name: Upload coverage to Codecov uses: codecov/codecov-action@v5 From 57d61a4fa91c9fb9ee142eae030aab13e373e469 Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:37:59 +0200 Subject: [PATCH 2/8] Adjust exceptions to start with AirOS --- airos/exceptions.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/airos/exceptions.py b/airos/exceptions.py index ac199e7..7708c01 100644 --- a/airos/exceptions.py +++ b/airos/exceptions.py @@ -5,33 +5,33 @@ class AirOSException(Exception): """Base error class for this AirOS library.""" -class ConnectionSetupError(AirOSException): +class AirOSConnectionSetupError(AirOSException): """Raised when unable to prepare authentication.""" -class ConnectionAuthenticationError(AirOSException): +class AirOSConnectionAuthenticationError(AirOSException): """Raised when unable to authenticate.""" -class DataMissingError(AirOSException): +class AirOSDataMissingError(AirOSException): """Raised when expected data is missing.""" -class KeyDataMissingError(AirOSException): +class AirOSKeyDataMissingError(AirOSException): """Raised when return data is missing critical keys.""" -class DeviceConnectionError(AirOSException): +class AirOSDeviceConnectionError(AirOSException): """Raised when unable to connect.""" -class AirosDiscoveryError(AirOSException): - """Base exception for Airos discovery issues.""" +class AirOSDiscoveryError(AirOSException): + """Base exception for AirOS discovery issues.""" -class AirosListenerError(AirosDiscoveryError): - """Raised when the Airos listener encounters an error.""" +class AirOSListenerError(AirOSDiscoveryError): + """Raised when the AirOS listener encounters an error.""" -class AirosEndpointError(AirosDiscoveryError): +class AirOSEndpointError(AirOSDiscoveryError): """Raised when there's an issue with the network endpoint.""" From e39dd78031b666cac5edce86bbfe16751caf4ef5 Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:38:19 +0200 Subject: [PATCH 3/8] Adjust to all caps --- airos/data.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airos/data.py b/airos/data.py index d969293..1de1d1f 100644 --- a/airos/data.py +++ b/airos/data.py @@ -42,8 +42,8 @@ class IeeeMode(Enum): class WirelessMode(Enum): """Enum definition.""" - AccessPoint_PointToPoint = "ap-ptp" - Station_PointToPoint = "sta-ptp" + PTP_ACCESSPOINT = "ap-ptp" + PTP_STATION = "sta-ptp" # More to be added when known From 6dbe600b40bc4d0a684707c413d71df99c9720cb Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:38:48 +0200 Subject: [PATCH 4/8] Add discovery for consumption + adhere new exceptions --- airos/discovery.py | 98 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 73 insertions(+), 25 deletions(-) diff --git a/airos/discovery.py b/airos/discovery.py index 5f85a05..512a691 100644 --- a/airos/discovery.py +++ b/airos/discovery.py @@ -7,7 +7,7 @@ import struct from typing import Any -from .exceptions import AirosDiscoveryError, AirosEndpointError, AirosListenerError +from .exceptions import AirOSDiscoveryError, AirOSEndpointError, AirOSListenerError _LOGGER = logging.getLogger(__name__) @@ -15,7 +15,7 @@ BUFFER_SIZE: int = 1024 -class AirosDiscoveryProtocol(asyncio.DatagramProtocol): +class AirOSDiscoveryProtocol(asyncio.DatagramProtocol): """A UDP protocol implementation for discovering Ubiquiti airOS devices. This class listens for UDP broadcast announcements from airOS devices @@ -30,7 +30,7 @@ class AirosDiscoveryProtocol(asyncio.DatagramProtocol): """ def __init__(self, callback: Callable[[dict[str, Any]], None]) -> None: - """Initialize AirosDiscoveryProtocol. + """Initialize AirOSDiscoveryProtocol. Args: callback: An asynchronous function to call when a device is discovered. @@ -46,7 +46,7 @@ def connection_made(self, transport: asyncio.BaseTransport) -> None: sock: socket.socket = self.transport.get_extra_info("socket") sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - log = f"Airos discovery listener (low-level) started on UDP port {DISCOVERY_PORT}." + log = f"AirOS discovery listener (low-level) started on UDP port {DISCOVERY_PORT}." _LOGGER.debug(log) def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None: @@ -60,30 +60,30 @@ def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None: if parsed_data: # Schedule the user-provided callback, don't await to keep listener responsive asyncio.create_task(self.callback(parsed_data)) # noqa: RUF006 - except (AirosEndpointError, AirosListenerError) as err: + except (AirOSEndpointError, AirOSListenerError) as err: # These are expected types of malformed packets. Log the specific error - # and then re-raise as AirosDiscoveryError. + # and then re-raise as AirOSDiscoveryError. log = f"Parsing failed for packet from {host_ip}: {err}" _LOGGER.exception(log) - raise AirosDiscoveryError(f"Malformed packet from {host_ip}") from err + raise AirOSDiscoveryError(f"Malformed packet from {host_ip}") from err except Exception as err: # General error during datagram reception (e.g., in callback itself) - log = f"Error processing Airos discovery packet from {host_ip}. Data hex: {data.hex()}" + log = f"Error processing AirOS discovery packet from {host_ip}. Data hex: {data.hex()}" _LOGGER.exception(log) - raise AirosDiscoveryError from err + raise AirOSDiscoveryError from err def error_received(self, exc: Exception | None) -> None: """Handle send or receive operation raises an OSError.""" if exc: - log = f"UDP error received in AirosDiscoveryProtocol: {exc}" + log = f"UDP error received in AirOSDiscoveryProtocol: {exc}" _LOGGER.error(log) def connection_lost(self, exc: Exception | None) -> None: """Handle connection is lost or closed.""" - _LOGGER.debug("AirosDiscoveryProtocol connection lost.") + _LOGGER.debug("AirOSDiscoveryProtocol connection lost.") if exc: - _LOGGER.exception("AirosDiscoveryProtocol connection lost due to") - raise AirosDiscoveryError from None + _LOGGER.exception("AirOSDiscoveryProtocol connection lost due to") + raise AirOSDiscoveryError from None def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None: """Parse a raw airOS discovery UDP packet. @@ -117,12 +117,12 @@ def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None if len(data) < 6: log = f"Packet too short for initial fixed header. Length: {len(data)}. Data: {data.hex()}" _LOGGER.debug(log) - raise AirosEndpointError(f"Malformed packet: {log}") + raise AirOSEndpointError(f"Malformed packet: {log}") if data[0] != 0x01 or data[1] != 0x06: - log = f"Packet does not start with expected Airos header (0x01 0x06). Actual: {data[0:2].hex()}" + log = f"Packet does not start with expected AirOS header (0x01 0x06). Actual: {data[0:2].hex()}" _LOGGER.debug(log) - raise AirosEndpointError(f"Malformed packet: {log}") + raise AirOSEndpointError(f"Malformed packet: {log}") offset: int = 6 @@ -151,7 +151,7 @@ def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None log = f"Truncated MAC address TLV (Type 0x06). Expected {expected_length}, got {len(data) - offset} bytes. Remaining: {data[offset:].hex()}" _LOGGER.warning(log) log = f"Malformed packet: {log}" - raise AirosEndpointError(log) + raise AirOSEndpointError(log) elif tlv_type in [ 0x02, @@ -169,7 +169,7 @@ def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None log = f"Truncated TLV (Type {tlv_type:#x}), no 2-byte length field. Remaining: {data[offset:].hex()}" _LOGGER.warning(log) log = f"Malformed packet: {log}" - raise AirosEndpointError(log) + raise AirOSEndpointError(log) tlv_length: int = struct.unpack_from(">H", data, offset)[0] offset += 2 @@ -182,7 +182,7 @@ def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None log = f"Data from TLV start: {data[offset - 3 :].hex()}" _LOGGER.warning(log) log = f"Malformed packet: {log}" - raise AirosEndpointError(log) + raise AirOSEndpointError(log) tlv_value: bytes = data[offset : offset + tlv_length] @@ -259,17 +259,65 @@ def parse_airos_packet(self, data: bytes, host_ip: str) -> dict[str, Any] | None log += f"Cannot determine length, stopping parsing. Remaining: {data[offset - 1 :].hex()}" _LOGGER.warning(log) log = f"Malformed packet: {log}" - raise AirosEndpointError(log) + raise AirOSEndpointError(log) except (struct.error, IndexError) as err: - log = f"Parsing error (struct/index) in AirosDiscoveryProtocol: {err} at offset {offset}. Remaining data: {data[offset:].hex()}" + log = f"Parsing error (struct/index) in AirOSDiscoveryProtocol: {err} at offset {offset}. Remaining data: {data[offset:].hex()}" _LOGGER.debug(log) log = f"Malformed packet: {log}" - raise AirosEndpointError(log) from err - except AirosEndpointError: # Catch AirosEndpointError specifically, re-raise it + raise AirOSEndpointError(log) from err + except AirOSEndpointError: # Catch AirOSEndpointError specifically, re-raise it raise except Exception as err: - _LOGGER.exception("Unexpected error during Airos packet parsing") - raise AirosListenerError from err + _LOGGER.exception("Unexpected error during AirOS packet parsing") + raise AirOSListenerError from err return parsed_info + + +async def async_discover_devices(timeout: int) -> dict[str, dict[str, Any]]: + """Discover unconfigured airOS devices on the network for a given timeout. + + This function sets up a listener, waits for a period, and returns + all discovered devices. + """ + _LOGGER.debug("Starting AirOS device discovery for %s seconds", timeout) + discovered_devices: dict[str, dict[str, Any]] = {} + + def _async_airos_device_found(device_info: dict[str, Any]) -> None: + """Handle discovered device.""" + mac_address = device_info.get("mac_address") + if mac_address: + discovered_devices[mac_address] = device_info + _LOGGER.debug("Discovered device: %s", device_info.get("hostname", mac_address)) + + transport: asyncio.DatagramTransport | None = None + try: + ( + transport, + protocol, + ) = await asyncio.get_running_loop().create_datagram_endpoint( + lambda: AirOSDiscoveryProtocol(_async_airos_device_found), + local_addr=("0.0.0.0", DISCOVERY_PORT), + ) + try: + await asyncio.sleep(timeout) + finally: + if transport: + _LOGGER.debug("Closing AirOS discovery listener") + transport.close() + except OSError as err: + if err.errno == 98: + _LOGGER.error("Address in use, another instance may be running.") + raise AirOSEndpointError("address_in_use") from err + _LOGGER.exception("Network endpoint error during discovery") + raise AirOSEndpointError("cannot_connect") from err + except asyncio.CancelledError as err: + _LOGGER.warning("Discovery listener cancelled: %s", err) + raise AirOSListenerError("cannot_connect") from err + except Exception as err: + _LOGGER.exception("An unexpected error occurred during discovery") + raise AirOSListenerError("cannot_connect") from err + + _LOGGER.debug("Discovery completed. Found %s devices.", len(discovered_devices)) + return discovered_devices From 850b5c75d1d9583c8359994862cd5dc5e387c7a5 Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:38:57 +0200 Subject: [PATCH 5/8] Adhere new exceptions --- airos/airos8.py | 95 ++++++++++++++++++++++++++++++++++--------------- 1 file changed, 67 insertions(+), 28 deletions(-) diff --git a/airos/airos8.py b/airos/airos8.py index f073fbf..a8915e0 100644 --- a/airos/airos8.py +++ b/airos/airos8.py @@ -12,11 +12,11 @@ from .data import AirOS8Data as AirOSData from .exceptions import ( - ConnectionAuthenticationError, - ConnectionSetupError, - DataMissingError, - DeviceConnectionError, - KeyDataMissingError, + AirOSConnectionAuthenticationError, + AirOSConnectionSetupError, + AirOSDataMissingError, + AirOSDeviceConnectionError, + AirOSKeyDataMissingError, ) _LOGGER = logging.getLogger(__name__) @@ -52,6 +52,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._provmode_url = f"{self.base_url}/api/provmode" # AirOS 8 self.current_csrf_token = None self._use_json_for_login_post = False @@ -103,10 +104,10 @@ async def login(self) -> bool: ) as response: if response.status == 403: _LOGGER.error("Authentication denied.") - raise ConnectionAuthenticationError from None + raise AirOSConnectionAuthenticationError from None if not response.cookies: _LOGGER.exception("Empty cookies after login, bailing out.") - raise ConnectionSetupError from None + raise AirOSConnectionSetupError from None else: for _, morsel in response.cookies.items(): # If the AIROS_ cookie was parsed but isn't automatically added to the jar, add it manually @@ -159,7 +160,7 @@ async def login(self) -> bool: _LOGGER.exception( "COOKIE JAR IS EMPTY after login POST. This is a major issue." ) - raise ConnectionSetupError from None + raise AirOSConnectionSetupError from None for cookie in self.session.cookie_jar: # pragma: no cover if cookie.key.startswith("AIROS_"): airos_cookie_found = True @@ -167,7 +168,7 @@ async def login(self) -> bool: ok_cookie_found = True if not airos_cookie_found and not ok_cookie_found: - raise ConnectionSetupError from None # pragma: no cover + raise AirOSConnectionSetupError from None # pragma: no cover response_text = await response.text() @@ -178,18 +179,18 @@ async def login(self) -> bool: return True except json.JSONDecodeError as err: _LOGGER.exception("JSON Decode Error") - raise DataMissingError from err + raise AirOSDataMissingError from err else: log = f"Login failed with status {response.status}. Full Response: {response.text}" _LOGGER.error(log) - raise ConnectionAuthenticationError from None + raise AirOSConnectionAuthenticationError from None except ( aiohttp.ClientError, aiohttp.client_exceptions.ConnectionTimeoutError, ) as err: _LOGGER.exception("Error during login") - raise DeviceConnectionError from err + raise AirOSDeviceConnectionError from err def derived_data( self, response: dict[str, Any] | None = None @@ -202,7 +203,7 @@ def derived_data( # No interfaces, no mac, no usability if not interfaces: - raise KeyDataMissingError from None + raise AirOSKeyDataMissingError from None for interface in interfaces: if interface["enabled"]: # Only consider if enabled @@ -227,7 +228,7 @@ async def status(self) -> AirOSData: """Retrieve status from the device.""" if not self.connected: _LOGGER.error("Not connected, login first") - raise DeviceConnectionError from None + raise AirOSDeviceConnectionError from None # --- Step 2: Verify authenticated access by fetching status.cgi --- authenticated_get_headers = {**self._common_headers} @@ -248,14 +249,14 @@ async def status(self) -> AirOSData: airos_data = AirOSData.from_dict(adjusted_json) except (MissingField, InvalidFieldValue) as err: _LOGGER.exception("Failed to deserialize AirOS data") - raise KeyDataMissingError from err + raise AirOSKeyDataMissingError from err return airos_data except json.JSONDecodeError: _LOGGER.exception( "JSON Decode Error in authenticated status response" ) - raise DataMissingError from None + raise AirOSDataMissingError from None else: log = f"Authenticated status.cgi failed: {response.status}. Response: {response_text}" _LOGGER.error(log) @@ -264,33 +265,32 @@ async def status(self) -> AirOSData: aiohttp.client_exceptions.ConnectionTimeoutError, ) as err: _LOGGER.exception("Error during authenticated status.cgi call") - raise DeviceConnectionError from err + raise AirOSDeviceConnectionError 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 + raise AirOSDeviceConnectionError from None if not mac_address: _LOGGER.error("Device mac-address missing") - raise DataMissingError from None + raise AirOSDataMissingError from None - kick_request_headers = {**self._common_headers} + request_headers = {**self._common_headers} if self.current_csrf_token: - kick_request_headers["X-CSRF-ID"] = self.current_csrf_token + request_headers["X-CSRF-ID"] = self.current_csrf_token - kick_payload = {"staif": "ath0", "staid": mac_address.upper()} + payload = {"staif": "ath0", "staid": mac_address.upper()} - kick_request_headers["Content-Type"] = ( + 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, + headers=request_headers, + data=payload, ) as response: if response.status == 200: return True @@ -302,5 +302,44 @@ async def stakick(self, mac_address: str = None) -> bool: aiohttp.ClientError, aiohttp.client_exceptions.ConnectionTimeoutError, ) as err: - _LOGGER.exception("Error during reconnect stakick.cgi call") - raise DeviceConnectionError from err + _LOGGER.exception("Error during reconnect request call") + raise AirOSDeviceConnectionError from err + + async def provmode(self, active: bool = False) -> bool: + """Set provisioning mode.""" + if not self.connected: + _LOGGER.error("Not connected, login first") + raise AirOSDeviceConnectionError from None + + request_headers = {**self._common_headers} + if self.current_csrf_token: + request_headers["X-CSRF-ID"] = self.current_csrf_token + + action = "stop" + if active: + action = "start" + + payload = {"action": action} + + request_headers["Content-Type"] = ( + "application/x-www-form-urlencoded; charset=UTF-8" + ) + + try: + async with self.session.post( + self._provmode_url, + headers=request_headers, + data=payload, + ) as response: + if response.status == 200: + return True + response_text = await response.text() + log = f"Unable to change provisioning mode response status {response.status} with {response_text}" + _LOGGER.error(log) + return False + except ( + aiohttp.ClientError, + aiohttp.client_exceptions.ConnectionTimeoutError, + ) as err: + _LOGGER.exception("Error during provisioning mode call") + raise AirOSDeviceConnectionError from err From 3799dac4d9357a364183722b3902011c1a661daa Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:39:19 +0200 Subject: [PATCH 6/8] Add appropriate testing --- tests/conftest.py | 26 ++++++++ tests/test_discovery.py | 130 +++++++++++++++++++++++++++++++++------- tests/test_stations.py | 8 +-- 3 files changed, 137 insertions(+), 27 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index d757c4c..b24ba09 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,12 @@ """Ubiquity AirOS test fixtures.""" +from unittest.mock import AsyncMock, MagicMock, patch from airos.airos8 import AirOS +from airos.discovery import AirOSDiscoveryProtocol import pytest import aiohttp +import asyncio @pytest.fixture @@ -19,3 +22,26 @@ async def airos_device(base_url): instance = AirOS(base_url, "username", "password", session, use_ssl=False) yield instance await session.close() + + +@pytest.fixture +def mock_datagram_endpoint(): + """Fixture to mock the creation of the UDP datagram endpoint.""" + # Define the mock objects FIRST, so they are in scope + mock_transport = MagicMock(spec=asyncio.DatagramTransport) + mock_protocol_instance = MagicMock(spec=AirOSDiscoveryProtocol) + + # Now, define the AsyncMock using the pre-defined variables + mock_create_datagram_endpoint = AsyncMock( + return_value=(mock_transport, mock_protocol_instance) + ) + + with patch( + "asyncio.get_running_loop" + ) as mock_get_loop, patch( + "airos.discovery.AirOSDiscoveryProtocol", new=MagicMock(return_value=mock_protocol_instance) + ): + mock_loop = mock_get_loop.return_value + mock_loop.create_datagram_endpoint = mock_create_datagram_endpoint + + yield mock_transport, mock_protocol_instance diff --git a/tests/test_discovery.py b/tests/test_discovery.py index 5d22e46..61ffdbb 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -5,8 +5,8 @@ import socket # Add this import from unittest.mock import AsyncMock, MagicMock, patch -from airos.discovery import DISCOVERY_PORT, AirosDiscoveryProtocol -from airos.exceptions import AirosDiscoveryError, AirosEndpointError +from airos.discovery import DISCOVERY_PORT, AirOSDiscoveryProtocol, async_discover_devices +from airos.exceptions import AirOSDiscoveryError, AirOSEndpointError, AirOSListenerError import pytest @@ -37,7 +37,7 @@ async def mock_airos_packet() -> bytes: @pytest.mark.asyncio async def test_parse_airos_packet_success(mock_airos_packet): """Test parse_airos_packet with a valid packet containing scrubbed data.""" - protocol = AirosDiscoveryProtocol( + protocol = AirOSDiscoveryProtocol( AsyncMock() ) # Callback won't be called directly in this unit test host_ip = ( @@ -61,17 +61,17 @@ async def test_parse_airos_packet_success(mock_airos_packet): @pytest.mark.asyncio async def test_parse_airos_packet_invalid_header(): """Test parse_airos_packet with an invalid header.""" - protocol = AirosDiscoveryProtocol(AsyncMock()) + protocol = AirOSDiscoveryProtocol(AsyncMock()) invalid_data = b"\x00\x00\x00\x00\x00\x00" + b"someotherdata" host_ip = "192.168.1.100" # Patch the _LOGGER.debug to verify the log message with patch("airos.discovery._LOGGER.debug") as mock_log_debug: - with pytest.raises(AirosEndpointError): + with pytest.raises(AirOSEndpointError): protocol.parse_airos_packet(invalid_data, host_ip) mock_log_debug.assert_called_once() assert ( - "does not start with expected Airos header" + "does not start with expected AirOS header" in mock_log_debug.call_args[0][0] ) @@ -79,13 +79,13 @@ async def test_parse_airos_packet_invalid_header(): @pytest.mark.asyncio async def test_parse_airos_packet_too_short(): """Test parse_airos_packet with data too short for header.""" - protocol = AirosDiscoveryProtocol(AsyncMock()) + protocol = AirOSDiscoveryProtocol(AsyncMock()) too_short_data = b"\x01\x06\x00" host_ip = "192.168.1.100" # Patch the _LOGGER.debug to verify the log message with patch("airos.discovery._LOGGER.debug") as mock_log_debug: - with pytest.raises(AirosEndpointError): + with pytest.raises(AirOSEndpointError): protocol.parse_airos_packet(too_short_data, host_ip) mock_log_debug.assert_called_once() assert ( @@ -97,7 +97,7 @@ async def test_parse_airos_packet_too_short(): @pytest.mark.asyncio async def test_parse_airos_packet_truncated_tlv(): """Test parse_airos_packet with a truncated TLV.""" - protocol = AirosDiscoveryProtocol(AsyncMock()) + protocol = AirOSDiscoveryProtocol(AsyncMock()) # Header + MAC TLV (valid) + then a truncated TLV_IP truncated_data = ( b"\x01\x06\x00\x00\x00\x00" # Header @@ -107,8 +107,8 @@ async def test_parse_airos_packet_truncated_tlv(): ) host_ip = "192.168.1.100" - # Expect AirosEndpointError due to struct.error or IndexError - with pytest.raises(AirosEndpointError): + # Expect AirOSEndpointError due to struct.error or IndexError + with pytest.raises(AirOSEndpointError): protocol.parse_airos_packet(truncated_data, host_ip) @@ -116,7 +116,7 @@ async def test_parse_airos_packet_truncated_tlv(): async def test_datagram_received_calls_callback(mock_airos_packet): """Test that datagram_received correctly calls the callback.""" mock_callback = AsyncMock() - protocol = AirosDiscoveryProtocol(mock_callback) + protocol = AirOSDiscoveryProtocol(mock_callback) host_ip = "192.168.1.3" # Sender IP with patch("asyncio.create_task") as mock_create_task: @@ -140,13 +140,13 @@ async def test_datagram_received_calls_callback(mock_airos_packet): async def test_datagram_received_handles_parsing_error(): """Test datagram_received handles exceptions during parsing.""" mock_callback = AsyncMock() - protocol = AirosDiscoveryProtocol(mock_callback) + protocol = AirOSDiscoveryProtocol(mock_callback) invalid_data = b"\x00\x00" # Too short, will cause parsing error host_ip = "192.168.1.100" with patch("airos.discovery._LOGGER.exception") as mock_log_exception: - # datagram_received catches errors internally and re-raises AirosDiscoveryError - with pytest.raises(AirosDiscoveryError): + # datagram_received catches errors internally and re-raises AirOSDiscoveryError + with pytest.raises(AirOSDiscoveryError): protocol.datagram_received(invalid_data, (host_ip, DISCOVERY_PORT)) mock_callback.assert_not_called() mock_log_exception.assert_called_once() # Ensure exception is logged @@ -155,7 +155,7 @@ async def test_datagram_received_handles_parsing_error(): @pytest.mark.asyncio async def test_connection_made_sets_transport(): """Test connection_made sets up transport and socket options.""" - protocol = AirosDiscoveryProtocol(AsyncMock()) + protocol = AirOSDiscoveryProtocol(AsyncMock()) mock_transport = MagicMock(spec=asyncio.DatagramTransport) mock_sock = MagicMock(spec=socket.socket) # Corrected: socket import added mock_transport.get_extra_info.return_value = mock_sock @@ -172,24 +172,24 @@ async def test_connection_made_sets_transport(): @pytest.mark.asyncio async def test_connection_lost_without_exception(): """Test connection_lost without an exception.""" - protocol = AirosDiscoveryProtocol(AsyncMock()) + protocol = AirOSDiscoveryProtocol(AsyncMock()) with patch("airos.discovery._LOGGER.debug") as mock_log_debug: protocol.connection_lost(None) mock_log_debug.assert_called_once_with( - "AirosDiscoveryProtocol connection lost." + "AirOSDiscoveryProtocol connection lost." ) @pytest.mark.asyncio async def test_connection_lost_with_exception(): """Test connection_lost with an exception.""" - protocol = AirosDiscoveryProtocol(AsyncMock()) + protocol = AirOSDiscoveryProtocol(AsyncMock()) test_exception = Exception("Test connection lost error") with ( patch("airos.discovery._LOGGER.exception") as mock_log_exception, pytest.raises( - AirosDiscoveryError - ), # connection_lost now re-raises AirosDiscoveryError + AirOSDiscoveryError + ), # connection_lost now re-raises AirOSDiscoveryError ): protocol.connection_lost(test_exception) mock_log_exception.assert_called_once() @@ -198,10 +198,94 @@ async def test_connection_lost_with_exception(): @pytest.mark.asyncio async def test_error_received(): """Test error_received logs the error.""" - protocol = AirosDiscoveryProtocol(AsyncMock()) + protocol = AirOSDiscoveryProtocol(AsyncMock()) test_exception = Exception("Test network error") with patch("airos.discovery._LOGGER.error") as mock_log_error: protocol.error_received(test_exception) mock_log_error.assert_called_once_with( - f"UDP error received in AirosDiscoveryProtocol: {test_exception}" + f"UDP error received in AirOSDiscoveryProtocol: {test_exception}" ) + +# Front-end discovery tests + +@pytest.mark.asyncio +async def test_async_discover_devices_success(mock_airos_packet, mock_datagram_endpoint): + """Test the high-level discovery function on a successful run.""" + mock_transport, mock_protocol_instance = mock_datagram_endpoint + + discovered_devices = {} + + def mock_protocol_factory(callback): + def inner_callback(device_info): + mac_address = device_info.get("mac_address") + if mac_address: + discovered_devices[mac_address] = device_info + + return MagicMock(callback=inner_callback) + + with patch( + "airos.discovery.AirOSDiscoveryProtocol", new=MagicMock(side_effect=mock_protocol_factory) + ): + + async def _simulate_discovery(): + await asyncio.sleep(0.1) + + protocol = AirOSDiscoveryProtocol(MagicMock()) # Create a real protocol instance just for parsing + parsed_data = protocol.parse_airos_packet(mock_airos_packet, "192.168.1.3") + + mock_protocol_factory(MagicMock()).callback(parsed_data) + + with patch("asyncio.sleep", new=AsyncMock()): + discovery_task = asyncio.create_task(async_discover_devices(timeout=1)) + + await _simulate_discovery() + + await discovery_task + + assert "01:23:45:67:89:CD" in discovered_devices + assert discovered_devices["01:23:45:67:89:CD"]["hostname"] == "name" + mock_transport.close.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_discover_devices_no_devices(mock_datagram_endpoint): + """Test discovery returns an empty dict if no devices are found.""" + mock_transport, _ = mock_datagram_endpoint + + with patch("asyncio.sleep", new=AsyncMock()): + result = await async_discover_devices(timeout=1) + + assert result == {} + mock_transport.close.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_discover_devices_oserror(mock_datagram_endpoint): + """Test discovery handles OSError during endpoint creation.""" + mock_transport, _ = mock_datagram_endpoint + + with patch( + "asyncio.get_running_loop" + ) as mock_get_loop, pytest.raises(AirOSEndpointError) as excinfo: + mock_loop = mock_get_loop.return_value + mock_loop.create_datagram_endpoint = AsyncMock(side_effect=OSError(98, "Address in use")) + + await async_discover_devices(timeout=1) + + assert "address_in_use" in str(excinfo.value) + mock_transport.close.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_discover_devices_cancelled(mock_datagram_endpoint): + """Test discovery handles CancelledError during the timeout.""" + mock_transport, _ = mock_datagram_endpoint + + # Patch asyncio.sleep to immediately raise CancelledError + with patch( + "asyncio.sleep", new=AsyncMock(side_effect=asyncio.CancelledError) + ), pytest.raises(AirOSListenerError) as excinfo: + await async_discover_devices(timeout=1) + + assert "cannot_connect" in str(excinfo.value) + mock_transport.close.assert_called_once() diff --git a/tests/test_stations.py b/tests/test_stations.py index a523b03..d7aa545 100644 --- a/tests/test_stations.py +++ b/tests/test_stations.py @@ -105,7 +105,7 @@ async def test_ap_corners(airos_device, base_url, mode="ap-ptp"): try: assert await airos_device.login() assert False - except airos.exceptions.ConnectionSetupError: + except airos.exceptions.AirOSConnectionSetupError: assert True mock_login_response.cookies = cookie @@ -124,7 +124,7 @@ async def test_ap_corners(airos_device, base_url, mode="ap-ptp"): try: assert await airos_device.login() assert False - except airos.exceptions.DataMissingError: + except airos.exceptions.AirOSDataMissingError: assert True mock_login_response.text = AsyncMock(return_value="{}") @@ -135,7 +135,7 @@ async def test_ap_corners(airos_device, base_url, mode="ap-ptp"): try: assert await airos_device.login() assert False - except airos.exceptions.ConnectionAuthenticationError: + except airos.exceptions.AirOSConnectionAuthenticationError: assert True mock_login_response.status = 200 @@ -143,5 +143,5 @@ async def test_ap_corners(airos_device, base_url, mode="ap-ptp"): try: assert await airos_device.login() assert False - except airos.exceptions.DeviceConnectionError: + except airos.exceptions.AirOSDeviceConnectionError: assert True From b8b8663837b755c38e606ebfbb0c4566a97605bd Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:39:29 +0200 Subject: [PATCH 7/8] Bump version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d597ad1..494b51a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "airos" -version = "0.2.1" +version = "0.2.2" license = "MIT" description = "Ubiquity airOS module(s) for Python 3." readme = "README.md" From 5ac774962b3f4e0091d0caece9d2a7aebab11eb7 Mon Sep 17 00:00:00 2001 From: Tom Scholten Date: Sat, 2 Aug 2025 11:39:54 +0200 Subject: [PATCH 8/8] Ruff --- airos/discovery.py | 4 +++- tests/conftest.py | 13 ++++++----- tests/test_discovery.py | 49 +++++++++++++++++++++++++++-------------- 3 files changed, 43 insertions(+), 23 deletions(-) diff --git a/airos/discovery.py b/airos/discovery.py index 512a691..93d772a 100644 --- a/airos/discovery.py +++ b/airos/discovery.py @@ -289,7 +289,9 @@ def _async_airos_device_found(device_info: dict[str, Any]) -> None: mac_address = device_info.get("mac_address") if mac_address: discovered_devices[mac_address] = device_info - _LOGGER.debug("Discovered device: %s", device_info.get("hostname", mac_address)) + _LOGGER.debug( + "Discovered device: %s", device_info.get("hostname", mac_address) + ) transport: asyncio.DatagramTransport | None = None try: diff --git a/tests/conftest.py b/tests/conftest.py index b24ba09..29aab81 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,12 +1,13 @@ """Ubiquity AirOS test fixtures.""" +import asyncio from unittest.mock import AsyncMock, MagicMock, patch + from airos.airos8 import AirOS from airos.discovery import AirOSDiscoveryProtocol import pytest import aiohttp -import asyncio @pytest.fixture @@ -36,10 +37,12 @@ def mock_datagram_endpoint(): return_value=(mock_transport, mock_protocol_instance) ) - with patch( - "asyncio.get_running_loop" - ) as mock_get_loop, patch( - "airos.discovery.AirOSDiscoveryProtocol", new=MagicMock(return_value=mock_protocol_instance) + with ( + patch("asyncio.get_running_loop") as mock_get_loop, + patch( + "airos.discovery.AirOSDiscoveryProtocol", + new=MagicMock(return_value=mock_protocol_instance), + ), ): mock_loop = mock_get_loop.return_value mock_loop.create_datagram_endpoint = mock_create_datagram_endpoint diff --git a/tests/test_discovery.py b/tests/test_discovery.py index 61ffdbb..dc142e8 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -5,7 +5,11 @@ import socket # Add this import from unittest.mock import AsyncMock, MagicMock, patch -from airos.discovery import DISCOVERY_PORT, AirOSDiscoveryProtocol, async_discover_devices +from airos.discovery import ( + DISCOVERY_PORT, + AirOSDiscoveryProtocol, + async_discover_devices, +) from airos.exceptions import AirOSDiscoveryError, AirOSEndpointError, AirOSListenerError import pytest @@ -206,10 +210,14 @@ async def test_error_received(): f"UDP error received in AirOSDiscoveryProtocol: {test_exception}" ) + # Front-end discovery tests + @pytest.mark.asyncio -async def test_async_discover_devices_success(mock_airos_packet, mock_datagram_endpoint): +async def test_async_discover_devices_success( + mock_airos_packet, mock_datagram_endpoint +): """Test the high-level discovery function on a successful run.""" mock_transport, mock_protocol_instance = mock_datagram_endpoint @@ -224,13 +232,16 @@ def inner_callback(device_info): return MagicMock(callback=inner_callback) with patch( - "airos.discovery.AirOSDiscoveryProtocol", new=MagicMock(side_effect=mock_protocol_factory) + "airos.discovery.AirOSDiscoveryProtocol", + new=MagicMock(side_effect=mock_protocol_factory), ): async def _simulate_discovery(): await asyncio.sleep(0.1) - protocol = AirOSDiscoveryProtocol(MagicMock()) # Create a real protocol instance just for parsing + protocol = AirOSDiscoveryProtocol( + MagicMock() + ) # Create a real protocol instance just for parsing parsed_data = protocol.parse_airos_packet(mock_airos_packet, "192.168.1.3") mock_protocol_factory(MagicMock()).callback(parsed_data) @@ -251,10 +262,10 @@ async def _simulate_discovery(): async def test_async_discover_devices_no_devices(mock_datagram_endpoint): """Test discovery returns an empty dict if no devices are found.""" mock_transport, _ = mock_datagram_endpoint - + with patch("asyncio.sleep", new=AsyncMock()): result = await async_discover_devices(timeout=1) - + assert result == {} mock_transport.close.assert_called_once() @@ -263,13 +274,16 @@ async def test_async_discover_devices_no_devices(mock_datagram_endpoint): async def test_async_discover_devices_oserror(mock_datagram_endpoint): """Test discovery handles OSError during endpoint creation.""" mock_transport, _ = mock_datagram_endpoint - - with patch( - "asyncio.get_running_loop" - ) as mock_get_loop, pytest.raises(AirOSEndpointError) as excinfo: + + with ( + patch("asyncio.get_running_loop") as mock_get_loop, + pytest.raises(AirOSEndpointError) as excinfo, + ): mock_loop = mock_get_loop.return_value - mock_loop.create_datagram_endpoint = AsyncMock(side_effect=OSError(98, "Address in use")) - + mock_loop.create_datagram_endpoint = AsyncMock( + side_effect=OSError(98, "Address in use") + ) + await async_discover_devices(timeout=1) assert "address_in_use" in str(excinfo.value) @@ -280,12 +294,13 @@ async def test_async_discover_devices_oserror(mock_datagram_endpoint): async def test_async_discover_devices_cancelled(mock_datagram_endpoint): """Test discovery handles CancelledError during the timeout.""" mock_transport, _ = mock_datagram_endpoint - + # Patch asyncio.sleep to immediately raise CancelledError - with patch( - "asyncio.sleep", new=AsyncMock(side_effect=asyncio.CancelledError) - ), pytest.raises(AirOSListenerError) as excinfo: + with ( + patch("asyncio.sleep", new=AsyncMock(side_effect=asyncio.CancelledError)), + pytest.raises(AirOSListenerError) as excinfo, + ): await async_discover_devices(timeout=1) - + assert "cannot_connect" in str(excinfo.value) mock_transport.close.assert_called_once()