Add appropriate testing

Author: CoMPaTech

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

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,31 +61,31 @@ 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]
         )
 
 
 @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,16 +107,16 @@ 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)
 
 
 @pytest.mark.asyncio
 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()

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,13 +135,13 @@ 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
     with patch.object(airos_device.session, "post", side_effect=aiohttp.ClientError):
         try:
             assert await airos_device.login()
             assert False
-        except airos.exceptions.DeviceConnectionError:
+        except airos.exceptions.AirOSDeviceConnectionError:
             assert True