From f00854656511e224b4222df0cadf86d9f56f74b4 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 24 Sep 2022 19:58:41 -0700
Subject: [PATCH] Avoid crashing if a device disconnects mid-diags

---
 scripts/wk/hw/diags.py |  5 +++++
 scripts/wk/hw/disk.py  | 11 +++++++++++
 scripts/wk/hw/smart.py |  5 +++++
 scripts/wk/os/linux.py |  4 ++++
 4 files changed, 25 insertions(+)

diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py
index 93980c37..855a608e 100644
--- a/scripts/wk/hw/diags.py
+++ b/scripts/wk/hw/diags.py
@@ -642,6 +642,11 @@ def disk_smart_status_check(dev, mid_run=True) -> None:
   color = None
   disable_tests = False
 
+  # Bail if dev is missing
+  if not dev.present:
+    dev.disable_disk_tests()
+    return
+
   # Check SMART status and attributes
   if not hw_smart.smart_status_ok(dev):
     msg = 'Critical SMART error detected'
diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py
index 45cbcce2..00d2f85e 100644
--- a/scripts/wk/hw/disk.py
+++ b/scripts/wk/hw/disk.py
@@ -151,12 +151,23 @@ class Disk:
 
     return aligned
 
+  @property
+  def present(self) -> bool:
+    """Verify this device is still present, returns bool."""
+    if not self.path.exists():
+      self.add_note('Device disconnected', 'RED')
+      return False
+    return True
+
   def update_details(self, skip_children=True) -> None:
     """Update disk details using OS specific methods.
 
     Required details default to generic descriptions
     and are converted to the correct type.
     """
+    if not self.present:
+      return
+
     if PLATFORM == 'Darwin':
       self.raw_details = get_disk_details_macos(
         self.path, skip_children=skip_children,
diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py
index f4292b06..927668d7 100644
--- a/scripts/wk/hw/smart.py
+++ b/scripts/wk/hw/smart.py
@@ -380,6 +380,11 @@ def update_smart_details(dev) -> None:
   """Update SMART details via smartctl."""
   updated_attributes = {}
 
+  # Bail if device was disconnected
+  if not dev.present:
+    dev.add_note('Device disconnected', 'RED')
+    return
+
   # Get SMART data
   cmd = [
     'sudo',
diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py
index f4aaa7eb..01b75c09 100644
--- a/scripts/wk/os/linux.py
+++ b/scripts/wk/os/linux.py
@@ -205,6 +205,10 @@ def mount_volumes(device_path=None, read_write=False, scan_corestorage=False):
     cmd.append(device_path)
   json_data = get_json_from_command(cmd)
 
+  # Bail if json_data is empty
+  if not json_data:
+    return
+
   # Build list of volumes
   for dev in json_data.get('blockdevices', [{}]):
     volumes.extend(_get_volumes(dev))