From b67080c462813879154629ccbbe3376a80568779 Mon Sep 17 00:00:00 2001
From: Nicholas Felt <nicholas.felt@tektronix.com>
Date: Mon, 11 Dec 2023 17:55:40 -0800
Subject: [PATCH] feat: Added a mechanism to reset cached properties whenever a
 device is rebooted. (#118)

---
 CHANGELOG.md                     |  4 ++++
 src/tm_devices/drivers/device.py | 30 ++++++++++++++++++++++--------
 tests/test_afgs.py               | 10 ++++++++++
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8b6f992d..ceab36c9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,10 @@ ______________________________________________________________________
 
 Things to be included in the next release go here.
 
+### Added
+
+- Added a step during a device reboot that will reset all the cached properties in the event that one of them changed.
+
 ### Changed
 
 - Switched to ruff's formatter instead of black's formatter for python code
diff --git a/src/tm_devices/drivers/device.py b/src/tm_devices/drivers/device.py
index 9b19a6ca..c66e1b56 100644
--- a/src/tm_devices/drivers/device.py
+++ b/src/tm_devices/drivers/device.py
@@ -3,7 +3,7 @@
 import time
 
 from abc import ABC, abstractmethod
-from contextlib import contextmanager
+from contextlib import contextmanager, suppress
 from functools import cached_property
 from typing import (
     Any,
@@ -89,13 +89,9 @@ def __str__(self) -> str:
         retval = f"{'=' * (line_break_length // 2)} {self.name} {'=' * (line_break_length // 2)}\n"
         retval += f"  {self.__class__} object at {id(self)}"
 
-        for prop in [
-            p
-            for p in dir(self.__class__)
-            if isinstance(getattr(self.__class__, p), (cached_property, property))
-            and not p.startswith("_")
-        ]:
-            retval += f"\n    {prop}={self.__getattribute__(prop)!r}"
+        for prop in self._get_self_properties():
+            if not prop.startswith("_"):
+                retval += f"\n    {prop}={self.__getattribute__(prop)!r}"
 
         retval += f"\n{'=' * (line_break_length + 2 + len(self.name))}"
         return retval
@@ -472,6 +468,16 @@ def reboot(self, quiet_period: int = 0) -> bool:
         Returns:
             A boolean representing the status of the reboot.
         """
+        # Reset the cached properties
+        for prop in self._get_self_properties():
+            if isinstance(getattr(self.__class__, prop), cached_property):
+                # Try to delete the cached_property, if it raises an AttributeError,
+                # that means that it has not previously been accessed and
+                # there is no need to delete the cached_property.
+                with suppress(AttributeError):
+                    self.__delattr__(prop)  # pylint: disable=unnecessary-dunder-call
+
+        # Reboot the device
         print_with_timestamp(f"Rebooting {self._name_and_alias}")
         self._reboot()
         self.close()
@@ -677,6 +683,14 @@ def wait_for_port_connection(
     # Private Methods
     ################################################################################################
 
+    def _get_self_properties(self) -> Tuple[str, ...]:
+        """Get a complete list of all the properties of the device."""
+        return tuple(
+            p
+            for p in dir(self.__class__)
+            if isinstance(getattr(self.__class__, p), (cached_property, property))
+        )
+
     @staticmethod
     @final
     def _verify_numerical_value(
diff --git a/tests/test_afgs.py b/tests/test_afgs.py
index d6917974..8cf16466 100644
--- a/tests/test_afgs.py
+++ b/tests/test_afgs.py
@@ -124,9 +124,19 @@ def test_afg31k(device_manager: DeviceManager, capsys: pytest.CaptureFixture[str
     afg31k = device_manager.add_afg("afg31k-hostname")
 
     _ = capsys.readouterr().out  # throw away stdout
+
+    # Check hostname
+    assert afg31k.hostname == "AFG31K-HOSTNAME"
+    # Change hostname to test the reset of cached properties
+    afg31k.hostname = "temp-hostname"
+    assert afg31k.hostname == "temp-hostname"
+
     # simulate a reboot
     afg31k.reboot()
 
+    # Test that the cached property was reset
+    assert afg31k.hostname == "AFG31K-HOSTNAME"
+
     stdout = capsys.readouterr().out
     assert "SYSTem:RESTart" in stdout