From d8bfd7b2f1d984fc469451494b1b57a45e1d772b Mon Sep 17 00:00:00 2001
From: pszulczewski <patryk@szulczewski.eu>
Date: Fri, 15 Dec 2023 16:46:07 +0100
Subject: [PATCH] Fix NXOS timeout

---
 pyntc/devices/nxos_device.py | 41 +++++++++++++-----------------------
 1 file changed, 15 insertions(+), 26 deletions(-)

diff --git a/pyntc/devices/nxos_device.py b/pyntc/devices/nxos_device.py
index a006739d..dbd54f37 100644
--- a/pyntc/devices/nxos_device.py
+++ b/pyntc/devices/nxos_device.py
@@ -7,7 +7,7 @@
 from pynxos.errors import CLIError
 from pynxos.features.file_copy import FileTransferError as NXOSFileTransferError
 
-from requests.exceptions import ReadTimeout
+from requests.exceptions import ReadTimeout, ConnectTimeout
 
 from pyntc import log
 from pyntc.devices.base_device import BaseDevice, fix_docs, RollbackError
@@ -53,22 +53,19 @@ def __init__(
 
     def _image_booted(self, image_name, **vendor_specifics):
         version_data = self.show("show version", raw_text=True)
-        if re.search(image_name, version_data):
-            log.info("Host %s: Image %s booted successfully.", self.host, image_name)
-            return True
-
-        log.info("Host %s: Image %s not booted successfully.", self.host, image_name)
-        return False
+        return bool(re.search(image_name, version_data))
 
     def _wait_for_device_reboot(self, timeout=3600):
         start = time.time()
         while time.time() - start < timeout:
+            self._uptime = None  # Reset cache
             try:
-                self.native.show("show hostname")
-                log.debug("Host %s: Device rebooted.", self.host)
-                return
+                # NXOS stays online, when it installs OS
+                if self.uptime < 180:
+                    log.debug("Host %s: Device rebooted.", self.host)
+                    return
             except:  # noqa E722 # nosec  # pylint: disable=bare-except
-                pass
+                time.sleep(10)
 
         log.error("Host %s: Device timed out while rebooting.", self.host)
         raise RebootTimeoutError(hostname=self.hostname, wait_time=timeout)
@@ -301,7 +298,9 @@ def install_os(self, image_name, **vendor_specifics):
         self.native.show("terminal dont-ask")
         timeout = vendor_specifics.get("timeout", 3600)
         if not self._image_booted(image_name):
+            log.info("Host %s: Setting Image %s in boot options.", self.host, image_name)
             self.set_boot_options(image_name, **vendor_specifics)
+            log.info("Host %s: Waiting for device reload.", self.host)
             self._wait_for_device_reboot(timeout=timeout)
             if not self._image_booted(image_name):
                 log.error("Host %s: OS install error for image %s", self.host, image_name)
@@ -311,7 +310,7 @@ def install_os(self, image_name, **vendor_specifics):
             log.info("Host %s: OS image %s installed successfully.", self.host, image_name)
             return True
 
-        log.info("Host %s: OS image %s not installed.", self.host, image_name)
+        log.info("Host %s: Image %s is already running on the device.", self.host, image_name)
         return False
 
     def open(self):  # noqa: D401
@@ -412,21 +411,11 @@ def set_boot_options(self, image_name, kickstart=None, **vendor_specifics):
             kickstart = file_system + kickstart
 
         image_name = file_system + image_name
-        # Allow for user defined timeout to take precedence if its over 300 seconds, otherwise change to 300.
         try:
-            native_timeout = int(self.native.timeout)
-        except (TypeError, ValueError):
-            native_timeout = 1
-
-        if native_timeout < 300:
-            self.native.timeout = 300
-        upgrade_result = self.native.set_boot_options(  # pylint: disable=assignment-from-no-return
-            image_name, kickstart=kickstart
-        )
-        self.native.timeout = 30
-
-        log.info("Host %s: boot options have been set to %s", self.host, upgrade_result)
-        return upgrade_result
+            self.native.set_boot_options(image_name, kickstart=kickstart)
+        except (ReadTimeout, ConnectTimeout):
+            pass
+        log.info("Host %s: boot options have been set to %s", self.host, image_name)
 
     def set_timeout(self, timeout):
         """Set timeout value on device connection.