diff options
Diffstat (limited to 'anta/device.py')
| -rw-r--r-- | anta/device.py | 43 |
1 files changed, 23 insertions, 20 deletions
diff --git a/anta/device.py b/anta/device.py index d517b8f..f0ec6a0 100644 --- a/anta/device.py +++ b/anta/device.py @@ -18,7 +18,8 @@ from aiocache.plugins import HitMissRatioPlugin from asyncssh import SSHClientConnection, SSHClientConnectionOptions from httpx import ConnectError, HTTPError, TimeoutException -from anta import __DEBUG__, aioeapi +import asynceapi +from anta import __DEBUG__ from anta.logger import anta_log_exception, exc_to_str from anta.models import AntaCommand @@ -116,7 +117,7 @@ class AntaDevice(ABC): yield "disable_cache", self.cache is None @abstractmethod - async def _collect(self, command: AntaCommand) -> None: + async def _collect(self, command: AntaCommand, *, collection_id: str | None = None) -> None: """Collect device command output. This abstract coroutine can be used to implement any command collection method @@ -131,11 +132,11 @@ class AntaDevice(ABC): Args: ---- - command: the command to collect - + command: The command to collect. + collection_id: An identifier used to build the eAPI request ID. """ - async def collect(self, command: AntaCommand) -> None: + async def collect(self, command: AntaCommand, *, collection_id: str | None = None) -> None: """Collect the output for a specified command. When caching is activated on both the device and the command, @@ -148,8 +149,8 @@ class AntaDevice(ABC): Args: ---- - command (AntaCommand): The command to process. - + command: The command to collect. + collection_id: An identifier used to build the eAPI request ID. """ # Need to ignore pylint no-member as Cache is a proxy class and pylint is not smart enough # https://github.com/pylint-dev/pylint/issues/7258 @@ -161,20 +162,20 @@ class AntaDevice(ABC): logger.debug("Cache hit for %s on %s", command.command, self.name) command.output = cached_output else: - await self._collect(command=command) + await self._collect(command=command, collection_id=collection_id) await self.cache.set(command.uid, command.output) # pylint: disable=no-member else: - await self._collect(command=command) + await self._collect(command=command, collection_id=collection_id) - async def collect_commands(self, commands: list[AntaCommand]) -> None: + async def collect_commands(self, commands: list[AntaCommand], *, collection_id: str | None = None) -> None: """Collect multiple commands. Args: ---- - commands: the commands to collect - + commands: The commands to collect. + collection_id: An identifier used to build the eAPI request ID. """ - await asyncio.gather(*(self.collect(command=command) for command in commands)) + await asyncio.gather(*(self.collect(command=command, collection_id=collection_id) for command in commands)) @abstractmethod async def refresh(self) -> None: @@ -270,7 +271,7 @@ class AsyncEOSDevice(AntaDevice): raise ValueError(message) self.enable = enable self._enable_password = enable_password - self._session: aioeapi.Device = aioeapi.Device(host=host, port=port, username=username, password=password, proto=proto, timeout=timeout) + self._session: asynceapi.Device = asynceapi.Device(host=host, port=port, username=username, password=password, proto=proto, timeout=timeout) ssh_params: dict[str, Any] = {} if insecure: ssh_params["known_hosts"] = None @@ -305,7 +306,7 @@ class AsyncEOSDevice(AntaDevice): """ return (self._session.host, self._session.port) - async def _collect(self, command: AntaCommand) -> None: # noqa: C901 function is too complex - because of many required except blocks + async def _collect(self, command: AntaCommand, *, collection_id: str | None = None) -> None: # noqa: C901 function is too complex - because of many required except blocks #pylint: disable=line-too-long """Collect device command output from EOS using aio-eapi. Supports outformat `json` and `text` as output structure. @@ -314,9 +315,10 @@ class AsyncEOSDevice(AntaDevice): Args: ---- - command: the AntaCommand to collect. + command: The command to collect. + collection_id: An identifier used to build the eAPI request ID. """ - commands: list[dict[str, Any]] = [] + commands: list[dict[str, str | int]] = [] if self.enable and self._enable_password is not None: commands.append( { @@ -329,14 +331,15 @@ class AsyncEOSDevice(AntaDevice): commands.append({"cmd": "enable"}) commands += [{"cmd": command.command, "revision": command.revision}] if command.revision else [{"cmd": command.command}] try: - response: list[dict[str, Any]] = await self._session.cli( + response: list[dict[str, Any] | str] = await self._session.cli( commands=commands, ofmt=command.ofmt, version=command.version, - ) + req_id=f"ANTA-{collection_id}-{id(command)}" if collection_id else f"ANTA-{id(command)}", + ) # type: ignore[assignment] # multiple commands returns a list # Do not keep response of 'enable' command command.output = response[-1] - except aioeapi.EapiCommandError as e: + except asynceapi.EapiCommandError as e: # This block catches exceptions related to EOS issuing an error. command.errors = e.errors if command.requires_privileges: |