diff --git a/CHANGELOG.md b/CHANGELOG.md index 23db4a44..e1bb74bc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,9 @@ -# unreleased +# Unreleased ## Performance - pyats broker: removed the per-command SSH liveness probe from the connection broker; dead sessions are now recovered by reconnecting and retrying the command once. Removes ~0.77s of event-loop blocking per test and the fleet-wide slowdown it caused. +- pyats broker: command rejections (SubCommandFailure) no longer trigger a full SSH disconnect and cache wipe; only transport-level failures do. Eliminates a 4.7x per-failure penalty on devices with unsupported commands. # 2.1.0b1 diff --git a/nac_test/pyats_core/broker/connection_broker.py b/nac_test/pyats_core/broker/connection_broker.py index c0099769..a68618c7 100644 --- a/nac_test/pyats_core/broker/connection_broker.py +++ b/nac_test/pyats_core/broker/connection_broker.py @@ -292,16 +292,29 @@ async def _run_and_cache(self, hostname: str, connection: Any, cmd: str) -> str: """Run a command on a connection once and cache its output. Raises: - Exception: Whatever the device layer raised, after tearing the - connection down so it is not handed to the next caller. + SubCommandFailure: Re-raised as-is. The device answered and rejected + the command, so the session and its cache are left intact. + Exception: Any other failure, after tearing the connection down so it + is not handed to the next caller. """ # Execute command in thread pool (since Unicon is synchronous) loop = get_or_create_event_loop() + try: + from unicon.core.errors import SubCommandFailure + except ImportError: + SubCommandFailure = None # type: ignore[assignment,misc] + try: output = await loop.run_in_executor(None, connection.execute, cmd) except Exception as e: + # SubCommandFailure is a fast path, not a health guarantee - unicon wraps + # transport errors in it too. Those get neither a disconnect here nor a retry + # in _execute_command; recovery relies on unicon re-establishing the session + # on the next execute(), which it does transparently. + if SubCommandFailure is not None and isinstance(e, SubCommandFailure): + logger.warning(f"Command rejected by {hostname} (session intact): {e}") + raise logger.error(f"Command execution failed on {hostname}: {e}") - # Try to reconnect on failure await self._disconnect_device(hostname) raise diff --git a/tests/integration/test_connection_broker_failure_modes.py b/tests/integration/test_connection_broker_failure_modes.py index 1d523fb3..57c79a26 100644 --- a/tests/integration/test_connection_broker_failure_modes.py +++ b/tests/integration/test_connection_broker_failure_modes.py @@ -27,6 +27,7 @@ from unittest.mock import MagicMock, patch import pytest +from unicon.core.errors import SubCommandFailure from nac_test.pyats_core.broker.broker_client import BrokerClient from nac_test.pyats_core.broker.connection_broker import ConnectionBroker @@ -680,3 +681,144 @@ async def _run() -> None: assert broker.connected_devices == {} assert not broker.socket_path.exists() + + +class TestCommandRejectionVsTransportFailure: + """Contrast SubCommandFailure (device rejects command, session healthy) with + transport-level failures (SSH dies, cache must be wiped). + + Both tests populate the command cache with a successful first command, then + trigger a failure on a second command and verify the divergent outcomes. + """ + + def test_sub_command_failure_preserves_connection_and_cache( + self, + make_broker: Any, + good_device: MagicMock, + ) -> None: + """SubCommandFailure re-raises to the client but does NOT disconnect + the device or wipe the command cache.""" + broker: ConnectionBroker = make_broker({"router-1": good_device}) + + async def _run() -> None: + loop = asyncio.get_event_loop() + + async def _body() -> None: + with patch( + "nac_test.pyats_core.broker.connection_broker.get_or_create_event_loop", + return_value=loop, + ): + with _patch_executor(loop): + async with BrokerClient( + socket_path=broker.socket_path + ) as client: + # 1) Successful command — populates cache + good_device.execute.return_value = "good output" + r1 = await asyncio.wait_for( + client._send_request( + { + "command": "execute", + "hostname": "router-1", + "cmd": "show version", + } + ), + timeout=2.0, + ) + assert r1["status"] == "success" + assert r1["result"] == "good output" + + # 2) SubCommandFailure on second command + good_device.execute.side_effect = SubCommandFailure( + "Invalid command at '^' marker" + ) + err = await _expect_broker_error( + client, + { + "command": "execute", + "hostname": "router-1", + "cmd": "show bgp all", + }, + timeout=2.0, + ) + assert "Invalid command" in err + + # Connection still present + assert "router-1" in broker.connected_devices + + # First command's cache entry still intact + cache = broker.command_cache.get("router-1") + assert cache is not None + assert cache.get("show version") == "good output" + + # Device was never disconnected + good_device.disconnect.assert_not_called() + + await _run_broker(broker, _body()) + + asyncio.run(_run()) + + def test_transport_failure_disconnects_and_wipes_cache( + self, + make_broker: Any, + good_device: MagicMock, + ) -> None: + """A transport-level exception (e.g. OSError) disconnects the device + and wipes the command cache — regression guard for existing behavior.""" + broker: ConnectionBroker = make_broker({"router-1": good_device}) + + async def _run() -> None: + loop = asyncio.get_event_loop() + + async def _body() -> None: + with patch( + "nac_test.pyats_core.broker.connection_broker.get_or_create_event_loop", + return_value=loop, + ): + with _patch_executor(loop): + async with BrokerClient( + socket_path=broker.socket_path + ) as client: + # 1) Successful command — populates cache + good_device.execute.return_value = "good output" + r1 = await asyncio.wait_for( + client._send_request( + { + "command": "execute", + "hostname": "router-1", + "cmd": "show version", + } + ), + timeout=2.0, + ) + assert r1["status"] == "success" + assert r1["result"] == "good output" + + # 2) Transport failure on second command + good_device.execute.side_effect = OSError( + "Socket is closed" + ) + err = await _expect_broker_error( + client, + { + "command": "execute", + "hostname": "router-1", + "cmd": "show interfaces", + }, + timeout=2.0, + ) + assert "Socket is closed" in err + + # OSError is a transport failure, so this exercises the full #899 path: + # fail -> disconnect -> reconnect -> fail -> disconnect -> raise + assert good_device.execute.call_count == 3 + assert good_device.disconnect.call_count == 2 + + # Connection removed + assert "router-1" not in broker.connected_devices + + # Command cache wiped for this device + assert "router-1" not in broker.command_cache + + await _run_broker(broker, _body()) + + asyncio.run(_run())