diff --git a/bellows/cli/util.py b/bellows/cli/util.py index 15ad66b2..2ff6942c 100644 --- a/bellows/cli/util.py +++ b/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): diff --git a/bellows/ezsp/__init__.py b/bellows/ezsp/__init__.py index df5bc13a..36566da4 100644 --- a/bellows/ezsp/__init__.py +++ b/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) diff --git a/tests/test_ezsp.py b/tests/test_ezsp.py index 211428ac..5d11b1af 100644 --- a/tests/test_ezsp.py +++ b/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(