diff --git a/tests/async_mock.py b/tests/async_mock.py
deleted file mode 100644
index 7d6b494..0000000
--- a/tests/async_mock.py
+++ /dev/null
@@ -1,7 +0,0 @@
-"""Mock utilities that are async aware."""
-import sys
-
-if sys.version_info[:2] < (3, 8):
-    from mock import *  # noqa
-else:
-    from unittest.mock import *  # noqa
diff --git a/tests/test_api.py b/tests/test_api.py
index 578c46f..cc8a4cb 100644
--- a/tests/test_api.py
+++ b/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
diff --git a/tests/test_application.py b/tests/test_application.py
index e34f92e..025981d 100644
--- a/tests/test_application.py
+++ b/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
\ No newline at end of file
diff --git a/tests/test_uart.py b/tests/test_uart.py
index 9d7b19a..1384b6a 100644
--- a/tests/test_uart.py
+++ b/tests/test_uart.py
@@ -1,4 +1,4 @@
-from .async_mock import MagicMock, AsyncMock
+from unittest.mock import MagicMock, AsyncMock
 
 import pytest
 import gpiozero
diff --git a/zigpy_zigate/zigbee/application.py b/zigpy_zigate/zigbee/application.py
index a903b2b..0686909 100644
--- a/zigpy_zigate/zigbee/application.py
+++ b/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')