Fix startup version command handling

tests/test_api.py

@@ -3,7 +3,7 @@
 import pytest
 import serial
 import serial_asyncio
-from .async_mock import AsyncMock, MagicMock, patch, sentinel
+from unittest.mock import AsyncMock, MagicMock, patch, sentinel
 
 import zigpy_zigate.config as config
 import zigpy_zigate.uart

tests/test_application.py

@@ -1,5 +1,4 @@
-from unittest import mock
-from .async_mock import AsyncMock, MagicMock, patch, sentinel
+from unittest.mock import AsyncMock, MagicMock, patch, sentinel
 
 import pytest
 import logging
@@ -28,7 +27,7 @@ def app():
 
 
 def test_zigpy_ieee(app):
-    cluster = mock.MagicMock()
+    cluster = MagicMock()
     cluster.cluster_id = 0x0000
     data = b"\x01\x02\x03\x04\x05\x06\x07\x08"
 
@@ -135,3 +134,18 @@ async def test_disconnect_multiple(app):
     await app.disconnect()
 
     assert app._api is None
+
+
+@pytest.mark.asyncio
+@patch("zigpy_zigate.zigbee.application.ZiGate.new")
+@pytest.mark.parametrize("version_rsp, expected_version", [
+    [((261, 798), 0), "3.1e"],
+    [((5, 801), 0), "3.21"]
+])
+async def test_startup_connect(zigate_new, app, version_rsp, expected_version):
+    api = zigate_new.return_value
+    api.version.return_value = version_rsp
+
+    await app.connect()
+
+    assert app.version == expected_version

tests/test_uart.py

@@ -1,4 +1,4 @@
-from .async_mock import MagicMock, AsyncMock
+from unittest.mock import MagicMock, AsyncMock
 
 import pytest
 import gpiozero

zigpy_zigate/zigbee/application.py

@@ -40,13 +40,13 @@ async def connect(self):
         api = await ZiGate.new(self._config[CONF_DEVICE], self)
         await api.set_raw_mode()
         await api.set_time()
-        version, lqi = await api.version()
-
-        self._api = api
 
+        (_, version), lqi = await api.version()
         major, minor = version.to_bytes(2, "big")
         self.version = f"{major:x}.{minor:x}"
 
+        self._api = api
+
         if self.version < '3.21':
             LOGGER.error('Old ZiGate firmware detected, you should upgrade to 3.21 or newer')