Fix radio probing broken by #485

bellows/cli/util.py

@@ -111,11 +111,12 @@ async def setup(dev, baudrate, cbh=None, configure=True):
         s.add_callback(cbh)
     try:
         await s.connect()
+        await s._startup_reset()
     except Exception as e:
         LOGGER.error(e)
+        s.close()
         raise click.Abort()
-    LOGGER.debug("Connected. Resetting.")
-    await s.reset()
+    LOGGER.debug("Connected")
     await s.version()
 
     async def cfg(config_id, value):

bellows/ezsp/__init__.py

@@ -74,31 +74,35 @@ async def probe(cls, device_config: Dict) -> bool | dict[str, int | str | bool]:
     async def _probe(self) -> None:
         """Open port and try sending a command"""
         await self.connect()
+        await self._startup_reset()
         await self.version()
 
-    @classmethod
-    async def initialize(cls, zigpy_config: Dict) -> "EZSP":
-        """Return initialized EZSP instance."""
-        ezsp = cls(zigpy_config[CONF_DEVICE])
-        await ezsp.connect()
-
+    async def _startup_reset(self):
+        """Start EZSP and reset the stack."""
         # `zigbeed` resets on startup
-        if zigpy_config[CONF_DEVICE][CONF_DEVICE_PATH].startswith("socket://"):
+        if self._config[CONF_DEVICE_PATH].startswith("socket://"):
             try:
                 await asyncio.wait_for(
-                    ezsp._gw.wait_for_startup_reset(),
+                    self._gw.wait_for_startup_reset(),
                     NETWORK_COORDINATOR_STARTUP_RESET_WAIT,
                 )
             except asyncio.TimeoutError:
                 pass
             else:
                 LOGGER.debug("Received a reset on startup, not resetting again")
-                ezsp.start_ezsp()
+                self.start_ezsp()
 
-        try:
-            if not ezsp.is_ezsp_running:
-                await ezsp.reset()
+        if not self.is_ezsp_running:
+            await self.reset()
 
+    @classmethod
+    async def initialize(cls, zigpy_config: Dict) -> "EZSP":
+        """Return initialized EZSP instance."""
+        ezsp = cls(zigpy_config[CONF_DEVICE])
+        await ezsp.connect()
+
+        try:
+            await ezsp._startup_reset()
             await ezsp.version()
             await ezsp._protocol.initialize(zigpy_config)
 

tests/test_ezsp.py

@@ -257,27 +257,34 @@ async def test_no_close_without_callback(ezsp_f):
     assert ezsp_f.close.call_count == 0
 
 
-@patch.object(ezsp.EZSP, "version", new_callable=AsyncMock)
 @patch("bellows.uart.connect", return_value=MagicMock(spec_set=uart.Gateway))
-async def test_probe_success(mock_connect, mock_version):
+async def test_probe_success(mock_connect):
     """Test device probing."""
 
-    res = await ezsp.EZSP.probe(DEVICE_CONFIG)
+    # Probe works with default baud
+    with patch(
+        "bellows.ezsp.protocol.ProtocolHandler.command",
+        AsyncMock(return_value=(4, 0, 0)),
+    ):
+        res = await ezsp.EZSP.probe(DEVICE_CONFIG)
+
     assert type(res) is dict
     assert mock_connect.call_count == 1
     assert mock_connect.await_count == 1
-    assert mock_version.call_count == 1
     assert mock_connect.return_value.close.call_count == 1
-
-    mock_connect.reset_mock()
-    mock_version.reset_mock()
     mock_connect.reset_mock()
-    res = await ezsp.EZSP.probe(DEVICE_CONFIG)
+
+    # Probe first fails with default baud but then works with alternate baud
+    with patch(
+        "bellows.ezsp.protocol.ProtocolHandler.command",
+        AsyncMock(side_effect=[asyncio.TimeoutError(), (4, 0, 0)]),
+    ):
+        res = await ezsp.EZSP.probe(DEVICE_CONFIG)
+
     assert type(res) is dict
-    assert mock_connect.call_count == 1
-    assert mock_connect.await_count == 1
-    assert mock_version.call_count == 1
-    assert mock_connect.return_value.close.call_count == 1
+    assert mock_connect.call_count == 2
+    assert mock_connect.await_count == 2
+    assert mock_connect.return_value.close.call_count == 2
 
 
 @pytest.mark.parametrize(