From 6f6d085a8bdbf31ff623102a7471a38599cf072b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20RAMAGE?= <sramage@poifindus.com>
Date: Fri, 20 Nov 2020 11:19:26 +0100
Subject: [PATCH] improve probing test

---
 tests/test_api.py    | 24 +++++++++++++++++-------
 zigpy_zigate/uart.py |  5 +++--
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/tests/test_api.py b/tests/test_api.py
index 9b65be5..cac4faa 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -1,5 +1,5 @@
 import asyncio
-
+import sys
 import pytest
 import serial
 import serial_asyncio
@@ -12,6 +12,7 @@
 DEVICE_CONFIG = config.SCHEMA_DEVICE({config.CONF_DEVICE_PATH: "/dev/null"})
 
 
+
 @pytest.fixture
 def api():
     api = zigate_api.ZiGate(DEVICE_CONFIG)
@@ -58,17 +59,26 @@ async def test_api_new(conn_mck):
 
 @pytest.mark.asyncio
 @patch.object(zigate_api.ZiGate, "set_raw_mode", new_callable=AsyncMock)
-@patch.object(zigpy_zigate.uart, "connect")
-async def test_probe_success(mock_connect, mock_raw_mode):
+@pytest.mark.parametrize(
+    "port",
+    ('/dev/null', 'pizigate:/dev/ttyAMA0'),
+)
+async def test_probe_success(mock_raw_mode, port, monkeypatch):
     """Test device probing."""
 
+    async def mock_conn(loop, protocol_factory, **kwargs):
+        protocol = protocol_factory()
+        loop.call_soon(protocol.connection_made, None)
+        return None, protocol
+    monkeypatch.setattr(serial_asyncio, "create_serial_connection", mock_conn)
+    DEVICE_CONFIG = zigpy_zigate.config.SCHEMA_DEVICE(
+        {zigpy_zigate.config.CONF_DEVICE_PATH: port}
+    )
+    sys.modules['RPi'] = MagicMock()
+    sys.modules['RPi.GPIO'] = MagicMock()
     res = await zigate_api.ZiGate.probe(DEVICE_CONFIG)
     assert res is True
-    assert mock_connect.call_count == 1
-    assert mock_connect.await_count == 1
-    assert mock_connect.call_args[0][0] == DEVICE_CONFIG
     assert mock_raw_mode.call_count == 1
-    assert mock_connect.return_value.close.call_count == 1
 
 
 @pytest.mark.asyncio
diff --git a/zigpy_zigate/uart.py b/zigpy_zigate/uart.py
index 9333c5b..fc0c3db 100644
--- a/zigpy_zigate/uart.py
+++ b/zigpy_zigate/uart.py
@@ -32,7 +32,8 @@ def connection_made(self, transport):
             self._connected_future.set_result(True)
 
     def close(self):
-        self._transport.close()
+        if self._transport:
+            self._transport.close()
 
     def send(self, cmd, data=b''):
         """Send data, taking care of escaping and framing"""
@@ -145,7 +146,7 @@ async def connect(device_config: Dict[str, Any], api, loop=None):
             # in case of pizigate:/dev/ttyAMA0 syntax
             if port.startswith('pizigate:'):
                 port = port[9:]
-        elif c.is_zigate_din:
+        elif c.is_zigate_din(port):
             LOGGER.debug('ZiGate USB DIN detected')
             await c.set_zigatedin_running_mode()