Fix UART not closing during probe

Author: puddly

tests/test_api.py

@@ -71,6 +71,9 @@ def pingable_serial_port(mocker):
     port_name = "/dev/ttyWorkingUSB1"
     old_serial_connect = serial_asyncio.create_serial_connection
 
+    def dummy_serial_close():
+        dummy_serial_conn.mock_transport._connected = False
+
     def dummy_serial_conn(loop, protocol_factory, url, *args, **kwargs):
         # Only our virtual port is handled differently
         if url != port_name:
@@ -84,6 +87,9 @@ def dummy_serial_conn(loop, protocol_factory, url, *args, **kwargs):
 
         fut.set_result((dummy_serial_conn.mock_transport, dummy_serial_conn.protocol))
 
+        assert not dummy_serial_conn.mock_transport._connected
+        dummy_serial_conn.mock_transport._connected = True
+
         return fut
 
     def ping_responder(data):
@@ -93,7 +99,9 @@ def ping_responder(data):
             dummy_serial_conn.protocol.data_received(b"\xFE\x02\x61\x01\x59\x06\x3D")
 
     dummy_serial_conn.mock_transport = mocker.Mock()
+    dummy_serial_conn.mock_transport._connected = False
     dummy_serial_conn.mock_transport.write = mocker.Mock(side_effect=ping_responder)
+    dummy_serial_conn.mock_transport.close = mocker.Mock(side_effect=dummy_serial_close)
 
     mocker.patch("serial_asyncio.create_serial_connection", new=dummy_serial_conn)
 

tests/test_application.py

@@ -726,18 +726,31 @@ async def test_on_af_message_callback(application, mocker):
 
 
 @pytest_mark_asyncio_timeout(seconds=3)
-async def test_probe(pingable_serial_port):  # noqa: F811
+async def test_probe_unsuccessful(pingable_serial_port):  # noqa: F811
     assert not (
         await ControllerApplication.probe(
             conf.SCHEMA_DEVICE({conf.CONF_DEVICE_PATH: "/dev/null"})
         )
     )
 
+
+@pytest_mark_asyncio_timeout(seconds=3)
+async def test_probe_successful(pingable_serial_port):  # noqa: F811
     assert await ControllerApplication.probe(
         conf.SCHEMA_DEVICE({conf.CONF_DEVICE_PATH: pingable_serial_port})
     )
 
 
+@pytest_mark_asyncio_timeout(seconds=3)
+async def test_probe_multiple(pingable_serial_port):  # noqa: F811
+    config = conf.SCHEMA_DEVICE({conf.CONF_DEVICE_PATH: pingable_serial_port})
+
+    assert await ControllerApplication.probe(config)
+    assert await ControllerApplication.probe(config)
+    assert await ControllerApplication.probe(config)
+    assert await ControllerApplication.probe(config)
+
+
 @pytest_mark_asyncio_timeout(seconds=5)
 async def test_reconnect(event_loop, application):
     app, znp_server = application()

zigpy_znp/api.py

@@ -220,7 +220,7 @@ async def connect(self, *, test_port=True, check_version=True) -> None:
                     else:
                         LOGGER.warning(old_version_msg)
         except Exception:
-            self._uart = None
+            self.close()
             raise
 
         LOGGER.debug(

zigpy_znp/zigbee/application.py

@@ -114,14 +114,16 @@ async def probe(cls, device_config: conf.ConfigType) -> bool:
 
         try:
             await znp.connect()
-            return True
+            result = True
         except Exception as e:
+            result = False
             LOGGER.warning(
                 "Failed to probe ZNP radio with config %s", device_config, exc_info=e
             )
-            return False
-        finally:
-            znp.close()
+
+        znp.close()
+
+        return result
 
     async def shutdown(self):
         """