Fix(anta.device): Make enable a flag for AsyncEOSDevice

aristanetworks · Jul 7, 2023 · 97e4272 · 97e4272
1 parent 0e9b74e
commit 97e4272
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/anta/device.py b/anta/device.py
@@ -172,6 +172,7 @@ def __init__(  # pylint: disable=R0913
         username: str,
         password: str,
         name: Optional[str] = None,
+        enable: bool = False,
         enable_password: Optional[str] = None,
         port: Optional[int] = None,
         ssh_port: Optional[int] = 22,
@@ -188,6 +189,7 @@ def __init__(  # pylint: disable=R0913
             username: Username to connect to eAPI and SSH
             password: Password to connect to eAPI and SSH
             name: Device name
+            enable: Device needs privileged access
             enable_password: Password used to gain privileged access on EOS
             proto: eAPI protocol. Value can be 'http' or 'https'
             port: eAPI port. Defaults to 80 is proto is 'http' or 443 if proto is 'https'.
@@ -199,6 +201,7 @@ def __init__(  # pylint: disable=R0913
         if name is None:
             name = f"{host}:{port}"
         super().__init__(name, tags)
+        self.enable = enable
         self._enable_password = enable_password
         self._session: Device = Device(host=host, port=port, username=username, password=password, proto=proto, timeout=timeout)
         ssh_params: Dict[str, Any] = {}
@@ -216,6 +219,7 @@ def __rich_repr__(self) -> Iterator[Tuple[str, Any]]:
         yield "eapi_port", self._session.port
         yield "username", self._ssh_opts.username
         yield "password", self._ssh_opts.password
+        yield "enable", self.enable
         yield "enable_password", self._enable_password
         yield "insecure", self._ssh_opts.known_hosts is None
         if __DEBUG__:
@@ -244,14 +248,15 @@ async def collect(self, command: AntaCommand) -> None:
         """
         try:
             commands = []
-            if self._enable_password is not None:
+            if self.enable and self._enable_password is not None:
                 commands.append(
                     {
                         "cmd": "enable",
                         "input": str(self._enable_password),
                     }
                 )
-            else:
+            elif self.enable:
+                # No password
                 commands.append({"cmd": "enable"})
             if command.revision:
                 commands.append({"cmd": command.command, "revision": command.revision})
@@ -266,7 +271,7 @@ async def collect(self, command: AntaCommand) -> None:
             # only applicable to json output
             if command.ofmt in ["json", "text"]:
                 # selecting only our command output
-                response = response[1]
+                response = response[-1]
             command.output = response
             logger.debug(f"{self.name}: {command}")