Handle port disconnects (#89)

* UART connection_lost() handler
* Handle UART disconnects.
* Update dependencies.

Author: Adminiuga

setup.py

@@ -13,6 +13,6 @@
     author_email="[email protected]",
     license="GPL-3.0",
     packages=find_packages(exclude=["*.tests"]),
-    install_requires=["pyserial-asyncio", "zigpy-homeassistant >= 0.10.0"],
+    install_requires=["pyserial-asyncio", "zigpy-homeassistant >= 0.17.0"],
     tests_require=["pytest"],
 )

tests/test_api.py

@@ -1,7 +1,9 @@
 import asyncio
-from unittest import mock
+import logging
 
+from asynctest import CoroutineMock, mock
 import pytest
+import zigpy.exceptions
 
 from zigpy_xbee import api as xbee_api, types as t, uart
 from zigpy_xbee.zigbee.application import ControllerApplication
@@ -25,9 +27,10 @@ async def test_connect(monkeypatch):
 
 
 def test_close(api):
-    api._uart.close = mock.MagicMock()
+    uart = api._uart
     api.close()
-    assert api._uart.close.call_count == 1
+    assert api._uart is None
+    assert uart.close.call_count == 1
 
 
 def test_commands():
@@ -84,6 +87,22 @@ def mock_api_frame(name, *args):
         api._uart.send.reset_mock()
 
 
+@pytest.mark.asyncio
+async def test_command_not_connected(api):
+    api._uart = None
+
+    def mock_api_frame(name, *args):
+        return mock.sentinel.api_frame_data, api._seq
+
+    api._api_frame = mock.MagicMock(side_effect=mock_api_frame)
+
+    for cmd, cmd_opts in xbee_api.COMMAND_REQUESTS.items():
+        with pytest.raises(zigpy.exceptions.APIException):
+            await api._command(cmd, mock.sentinel.cmd_data)
+        assert api._api_frame.call_count == 0
+        api._api_frame.reset_mock()
+
+
 async def _test_at_or_queued_at_command(api, cmd, monkeypatch, do_reply=True):
     monkeypatch.setattr(
         t, "serialize", mock.MagicMock(return_value=mock.sentinel.serialize)
@@ -518,3 +537,54 @@ def test_handle_many_to_one_rri(api):
     ieee = t.EUI64([t.uint8_t(a) for a in range(0, 8)])
     nwk = 0x1234
     api._handle_many_to_one_rri(ieee, nwk, 0)
+
+
+@pytest.mark.asyncio
+async def test_reconnect_multiple_disconnects(monkeypatch, caplog):
+    api = xbee_api.XBee()
+    dev = mock.sentinel.uart
+    connect_mock = CoroutineMock()
+    connect_mock.return_value = asyncio.Future()
+    connect_mock.return_value.set_result(True)
+    monkeypatch.setattr(uart, "connect", connect_mock)
+
+    await api.connect(dev, 115200)
+
+    caplog.set_level(logging.DEBUG)
+    connected = asyncio.Future()
+    connected.set_result(mock.sentinel.uart_reconnect)
+    connect_mock.reset_mock()
+    connect_mock.side_effect = [asyncio.Future(), connected]
+    api.connection_lost("connection lost")
+    await asyncio.sleep(0.3)
+    api.connection_lost("connection lost 2")
+    await asyncio.sleep(0.3)
+
+    assert "Cancelling reconnection attempt" in caplog.messages
+    assert api._uart is mock.sentinel.uart_reconnect
+    assert connect_mock.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_reconnect_multiple_attempts(monkeypatch, caplog):
+    api = xbee_api.XBee()
+    dev = mock.sentinel.uart
+    connect_mock = CoroutineMock()
+    connect_mock.return_value = asyncio.Future()
+    connect_mock.return_value.set_result(True)
+    monkeypatch.setattr(uart, "connect", connect_mock)
+
+    await api.connect(dev, 115200)
+
+    caplog.set_level(logging.DEBUG)
+    connected = asyncio.Future()
+    connected.set_result(mock.sentinel.uart_reconnect)
+    connect_mock.reset_mock()
+    connect_mock.side_effect = [asyncio.TimeoutError, OSError, connected]
+
+    with mock.patch("asyncio.sleep"):
+        api.connection_lost("connection lost")
+        await api._conn_lost_task
+
+    assert api._uart is mock.sentinel.uart_reconnect
+    assert connect_mock.call_count == 3

tests/test_uart.py

@@ -1,3 +1,4 @@
+import asyncio
 from unittest import mock
 
 import pytest
@@ -202,3 +203,24 @@ def test_unescape_underflow(gw):
     unescaped, rest = gw._get_unescaped(escaped, 3)
     assert unescaped is None
     assert rest is None
+
+
+def test_connection_lost_exc(gw):
+    gw._connected_future = asyncio.Future()
+
+    gw.connection_lost(ValueError())
+
+    conn_lost = gw._api.connection_lost
+    assert conn_lost.call_count == 1
+    assert isinstance(conn_lost.call_args[0][0], Exception)
+    assert gw._connected_future.done()
+    assert gw._connected_future.exception()
+
+
+def test_connection_closed(gw):
+    gw._connected_future = asyncio.Future()
+    gw.connection_lost(None)
+
+    assert gw._api.connection_lost.call_count == 0
+    assert gw._connected_future.done()
+    assert gw._connected_future.result() is True

tox.ini

@@ -12,6 +12,7 @@ setenv = PYTHONPATH = {toxinidir}
 install_command = pip install {opts} {packages}
 commands = py.test --cov --cov-report=
 deps =
+    asynctest
     coveralls
     pytest
     pytest-cov

zigpy_xbee/api.py

@@ -4,7 +4,7 @@
 import functools
 import logging
 
-from zigpy.exceptions import DeliveryError
+from zigpy.exceptions import APIException, DeliveryError
 from zigpy.types import LVList
 
 from . import types as t, uart
@@ -243,11 +243,13 @@ class ATCommandResult(enum.IntEnum):
 class XBee:
     def __init__(self):
         self._uart = None
+        self._uart_params = None
         self._seq = 1
         self._commands_by_id = {v[0]: k for k, v in COMMAND_RESPONSES.items()}
         self._awaiting = {}
         self._app = None
         self._cmd_mode_future = None
+        self._conn_lost_task = None
         self._reset = asyncio.Event()
         self._running = asyncio.Event()
 
@@ -266,15 +268,69 @@ def is_running(self):
         """Return true if coordinator is running."""
         return self.coordinator_started_event.is_set()
 
-    async def connect(self, device, baudrate=115200):
+    async def connect(self, device: str, baudrate: int = 115200) -> None:
         assert self._uart is None
         self._uart = await uart.connect(device, baudrate, self)
+        self._uart_params = (device, baudrate)
+
+    def reconnect(self):
+        """Reconnect using saved parameters."""
+        LOGGER.debug(
+            "Reconnecting '%s' serial port using %s",
+            self._uart_params[0],
+            self._uart_params[1],
+        )
+        return self.connect(self._uart_params[0], self._uart_params[1])
+
+    def connection_lost(self, exc: Exception) -> None:
+        """Lost serial connection."""
+        LOGGER.warning(
+            "Serial '%s' connection lost unexpectedly: %s", self._uart_params[0], exc
+        )
+        self._uart = None
+        if self._conn_lost_task and not self._conn_lost_task.done():
+            self._conn_lost_task.cancel()
+        self._conn_lost_task = asyncio.ensure_future(self._connection_lost())
+
+    async def _connection_lost(self) -> None:
+        """Reconnect serial port."""
+        try:
+            await self._reconnect_till_done()
+        except asyncio.CancelledError:
+            LOGGER.debug("Cancelling reconnection attempt")
+
+    async def _reconnect_till_done(self) -> None:
+        attempt = 1
+        while True:
+            try:
+                await asyncio.wait_for(self.reconnect(), timeout=10)
+                break
+            except (asyncio.TimeoutError, OSError) as exc:
+                wait = 2 ** min(attempt, 5)
+                attempt += 1
+                LOGGER.debug(
+                    "Couldn't re-open '%s' serial port, retrying in %ss: %s",
+                    self._uart_params[0],
+                    wait,
+                    str(exc),
+                )
+                await asyncio.sleep(wait)
+
+        LOGGER.debug(
+            "Reconnected '%s' serial port after %s attempts",
+            self._uart_params[0],
+            attempt,
+        )
 
     def close(self):
-        return self._uart.close()
+        if self._uart:
+            self._uart.close()
+            self._uart = None
 
     def _command(self, name, *args, mask_frame_id=False):
         LOGGER.debug("Command %s %s", name, args)
+        if self._uart is None:
+            raise APIException("API is not running")
         frame_id = 0 if mask_frame_id else self._seq
         data, needs_response = self._api_frame(name, frame_id, *args)
         self._uart.send(data)

zigpy_xbee/uart.py

@@ -14,6 +14,7 @@ class Gateway(asyncio.Protocol):
     XOFF = b"\x13"
 
     RESERVED = START + ESCAPE + XON + XOFF
+    THIS_ONE = True
 
     def __init__(self, api, connected_future=None):
         self._buffer = b""
@@ -45,6 +46,20 @@ def baudrate(self, baudrate):
                 "baudrate must be one of {}".format(self._transport.serial.BAUDRATES)
             )
 
+    def connection_lost(self, exc) -> None:
+        """Port was closed expectedly or unexpectedly."""
+        if self._connected_future and not self._connected_future.done():
+            if exc is None:
+                self._connected_future.set_result(True)
+            else:
+                self._connected_future.set_exception(exc)
+        if exc is None:
+            LOGGER.debug("Closed serial connection")
+            return
+
+        LOGGER.error("Lost serial connection: %s", exc)
+        self._api.connection_lost(exc)
+
     def connection_made(self, transport):
         """Callback when the uart is connected"""
         LOGGER.debug("Connection made")