Implement radio probe (#90)

* Probe method.
* Relax command mode AT response decoding.
* Update tests.

Author: Adminiuga

tests/test_api.py

@@ -3,6 +3,7 @@
 
 from asynctest import CoroutineMock, mock
 import pytest
+import serial
 import zigpy.exceptions
 
 from zigpy_xbee import api as xbee_api, types as t, uart
@@ -588,3 +589,80 @@ async def test_reconnect_multiple_attempts(monkeypatch, caplog):
 
     assert api._uart is mock.sentinel.uart_reconnect
     assert connect_mock.call_count == 3
+
+
+@pytest.mark.asyncio
+@mock.patch.object(xbee_api.XBee, "_at_command", new_callable=CoroutineMock)
+@mock.patch.object(uart, "connect")
+async def test_probe_success(mock_connect, mock_at_cmd):
+    """Test device probing."""
+
+    res = await xbee_api.XBee.probe(mock.sentinel.uart, mock.sentinel.baud)
+    assert res is True
+    assert mock_connect.call_count == 1
+    assert mock_connect.await_count == 1
+    assert mock_connect.call_args[0][0] is mock.sentinel.uart
+    assert mock_at_cmd.call_count == 1
+    assert mock_connect.return_value.close.call_count == 1
+
+
+@pytest.mark.asyncio
+@mock.patch.object(xbee_api.XBee, "init_api_mode", return_value=True)
+@mock.patch.object(xbee_api.XBee, "_at_command", side_effect=asyncio.TimeoutError)
+@mock.patch.object(uart, "connect")
+async def test_probe_success_api_mode(mock_connect, mock_at_cmd, mock_api_mode):
+    """Test device probing."""
+
+    res = await xbee_api.XBee.probe(mock.sentinel.uart, mock.sentinel.baud)
+    assert res is True
+    assert mock_connect.call_count == 1
+    assert mock_connect.await_count == 1
+    assert mock_connect.call_args[0][0] is mock.sentinel.uart
+    assert mock_at_cmd.call_count == 1
+    assert mock_api_mode.call_count == 1
+    assert mock_connect.return_value.close.call_count == 1
+
+
+@pytest.mark.asyncio
+@mock.patch.object(xbee_api.XBee, "init_api_mode")
+@mock.patch.object(xbee_api.XBee, "_at_command", side_effect=asyncio.TimeoutError)
+@mock.patch.object(uart, "connect")
+@pytest.mark.parametrize(
+    "exception",
+    (asyncio.TimeoutError, serial.SerialException, zigpy.exceptions.APIException),
+)
+async def test_probe_fail(mock_connect, mock_at_cmd, mock_api_mode, exception):
+    """Test device probing fails."""
+
+    mock_api_mode.side_effect = exception
+    mock_api_mode.reset_mock()
+    mock_at_cmd.reset_mock()
+    mock_connect.reset_mock()
+    res = await xbee_api.XBee.probe(mock.sentinel.uart, mock.sentinel.baud)
+    assert res is False
+    assert mock_connect.call_count == 1
+    assert mock_connect.await_count == 1
+    assert mock_connect.call_args[0][0] is mock.sentinel.uart
+    assert mock_at_cmd.call_count == 1
+    assert mock_api_mode.call_count == 1
+    assert mock_connect.return_value.close.call_count == 1
+
+
+@pytest.mark.asyncio
+@mock.patch.object(xbee_api.XBee, "init_api_mode", return_value=False)
+@mock.patch.object(xbee_api.XBee, "_at_command", side_effect=asyncio.TimeoutError)
+@mock.patch.object(uart, "connect")
+async def test_probe_fail_api_mode(mock_connect, mock_at_cmd, mock_api_mode):
+    """Test device probing fails."""
+
+    mock_api_mode.reset_mock()
+    mock_at_cmd.reset_mock()
+    mock_connect.reset_mock()
+    res = await xbee_api.XBee.probe(mock.sentinel.uart, mock.sentinel.baud)
+    assert res is False
+    assert mock_connect.call_count == 1
+    assert mock_connect.await_count == 1
+    assert mock_connect.call_args[0][0] is mock.sentinel.uart
+    assert mock_at_cmd.call_count == 1
+    assert mock_api_mode.call_count == 1
+    assert mock_connect.return_value.close.call_count == 1

tests/test_uart.py

@@ -48,13 +48,6 @@ def test_command_mode_rsp(gw):
     assert gw._api.handle_command_mode_rsp.call_args[0][0] == "OK"
 
 
-def test_command_mode_rsp_decode_exc(gw):
-    data = b"OK\x81"
-    with pytest.raises(UnicodeDecodeError):
-        gw.command_mode_rsp(data)
-    assert gw._api.handle_command_mode_rsp.call_count == 0
-
-
 def test_command_mode_send(gw):
     data = b"ATAP2\x0D"
     gw.command_mode_send(data)

zigpy_xbee/api.py

@@ -4,6 +4,7 @@
 import functools
 import logging
 
+import serial
 from zigpy.exceptions import APIException, DeliveryError
 from zigpy.types import LVList
 
@@ -13,6 +14,7 @@
 
 AT_COMMAND_TIMEOUT = 2
 REMOTE_AT_COMMAND_TIMEOUT = 30
+PROBE_TIMEOUT = 45
 
 
 class ModemStatus(t.uint8_t, t.UndefinedEnum):
@@ -480,9 +482,9 @@ def handle_command_mode_rsp(self, data):
         fut = self._cmd_mode_future
         if fut is None or fut.done():
             return
-        if data == "OK":
+        if "OK" in data:
             fut.set_result(True)
-        elif data == "ERROR":
+        elif "ERROR" in data:
             fut.set_result(False)
         else:
             fut.set_result(data)
@@ -548,6 +550,32 @@ async def init_api_mode(self):
         )
         return False
 
+    @classmethod
+    async def probe(cls, device: str, baudrate: int) -> bool:
+        """Probe port for the device presence."""
+        api = cls()
+        try:
+            await asyncio.wait_for(api._probe(device, baudrate), timeout=PROBE_TIMEOUT)
+            return True
+        except (asyncio.TimeoutError, serial.SerialException, APIException) as exc:
+            LOGGER.debug("Unsuccessful radio probe of '%s' port", exc_info=exc)
+        finally:
+            api.close()
+
+        return False
+
+    async def _probe(self, device: str, baudrate: int) -> None:
+        """Open port and try sending a command"""
+        await self.connect(device, baudrate)
+        try:
+            # Ensure we have escaped commands
+            await self._at_command("AP", 2)
+        except asyncio.TimeoutError:
+            if not await self.init_api_mode():
+                raise APIException("Failed to configure XBee for API mode")
+        finally:
+            self.close()
+
     def __getattr__(self, item):
         if item in COMMAND_REQUESTS:
             return functools.partial(self._command, item)

zigpy_xbee/uart.py

@@ -69,11 +69,7 @@ def connection_made(self, transport):
 
     def command_mode_rsp(self, data):
         """Handles AT command mode response."""
-        try:
-            data = data.decode("ascii")
-        except UnicodeDecodeError as ex:
-            LOGGER.debug("Couldn't ascii decode AT command mode response: %s", ex)
-            raise
+        data = data.decode("ascii", "ignore")
         LOGGER.debug("Handling AT command mode response: %s", data)
         self._api.handle_command_mode_rsp(data)