From 2706d1a9a5a438fd63bde1571b2e8cbc5b0de2aa Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 23 Oct 2022 13:18:43 -0700
Subject: [PATCH] Show failed SMART attributes during disk tests

---
 scripts/wk/hw/diags.py | 25 ++++++++--------
 scripts/wk/hw/smart.py | 67 ++++++++++++++++++++++++------------------
 2 files changed, 51 insertions(+), 41 deletions(-)

diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py
index 855a608e..2c1fa61b 100644
--- a/scripts/wk/hw/diags.py
+++ b/scripts/wk/hw/diags.py
@@ -594,6 +594,7 @@ def disk_self_test(state, test_objects, test_mode=False) -> None:
     f'Disk self-test{"s" if len(test_objects) > 1 else ""}',
     )
   std.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}')
+  show_failed_attributes(state)
   for test in reversed(test_objects):
     if test.disabled:
       # Skip
@@ -692,19 +693,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None:
   std.print_info(
     f'Starting disk surface scan{"s" if len(test_objects) > 1 else ""}',
     )
-  for disk in state.disks:
-    failed_attributes = [
-      line for line in hw_smart.generate_attribute_report(disk) if 'failed' in line
-      ]
-    if failed_attributes:
-      size_str = std.bytes_to_string(disk.size, use_binary=False)
-      std.print_colored(
-        ['[', disk.path.name, ' ', size_str, ']'],
-        [None, 'BLUE', None, 'CYAN', None],
-        sep='',
-        )
-      std.print_report(failed_attributes)
-      std.print_standard('')
+  show_failed_attributes(state)
 
   # Run surface scans
   for test in reversed([test for test in test_objects if not test.disabled]):
@@ -907,6 +896,16 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None:
     std.pause('Press Enter to return to main menu...')
 
 
+def show_failed_attributes(state) -> None:
+  """Show failed attributes for all disks."""
+  for dev in state.disks:
+    std.print_colored([dev.name, dev.description], ['CYAN', None])
+    std.print_report(
+      hw_smart.generate_attribute_report(dev, only_failed=True),
+      )
+    std.print_standard('')
+
+
 def show_results(state) -> None:
   """Show test results by device."""
   std.sleep(0.5)
diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py
index 85afebe3..d9f47f7a 100644
--- a/scripts/wk/hw/smart.py
+++ b/scripts/wk/hw/smart.py
@@ -112,20 +112,14 @@ def enable_smart(dev) -> None:
   run_program(cmd, check=False)
 
 
-def generate_attribute_report(dev) -> list[str]:
+def generate_attribute_report(dev, only_failed=False) -> list[str]:
   """Generate attribute report, returns list."""
   report = []
   for attr, value in sorted(dev.attributes.items()):
-    note = ''
-    value_color = 'GREEN'
-
     # Skip attributes not in our list
     if attr not in dev.known_attributes:
       continue
 
-    # Check for attribute note
-    note = dev.known_attributes[attr].get('Note', '')
-
     # ID / Name
     label = f'{attr:>3}'
     if isinstance(attr, int):
@@ -133,28 +127,13 @@ def generate_attribute_report(dev) -> list[str]:
       label += f' / {str(hex(attr))[2:].upper():0>2}: {value["name"]}'
     label = f'  {label.replace("_", " "):38}'
 
-    # Value color
-    if dev.known_attributes[attr].get('PercentageLife', False):
-      # PercentageLife values
-      if 0 <= value['raw'] <= dev.known_attributes[attr]['Error']:
-        value_color = 'RED'
-        note = '(failed, % life remaining)'
-      elif value['raw'] < 0 or value['raw'] > 100:
-        value_color = 'PURPLE'
-        note = '(invalid?)'
-    else:
-      for threshold, color in ATTRIBUTE_COLORS:
-        threshold_val = dev.known_attributes[attr].get(threshold, None)
-        if threshold_val and value['raw'] >= threshold_val:
-          value_color = color
-          if threshold == 'Error':
-            note = '(failed)'
-          elif threshold == 'Maximum':
-            note = '(invalid?)'
+    # Color & Note
+    value_color, note = get_attribute_value_color_and_note(dev, attr, value)
 
-    # 199/C7 warning
-    if str(attr) == '199' and value['raw'] > 0:
-      note = '(bad cable?)'
+    # Skip non-failing attributes if requested
+    ## NOTE: This is a naive test and will include 'invalid' attributes
+    if only_failed and not note:
+      continue
 
     # Build colored string and append to report
     line = color_string(
@@ -167,6 +146,38 @@ def generate_attribute_report(dev) -> list[str]:
   return report
 
 
+def get_attribute_value_color_and_note(dev, attr, value) -> tuple[str, str]:
+  """Get attribute color and note based on SMART data."""
+  value_color = 'GREEN'
+  note = dev.known_attributes[attr].get('Note', '')
+
+  # Value value_color
+  if dev.known_attributes[attr].get('PercentageLife', False):
+    # PercentageLife values
+    if 0 <= value['raw'] <= dev.known_attributes[attr]['Error']:
+      value_color = 'RED'
+      note = '(failed, % life remaining)'
+    elif value['raw'] < 0 or value['raw'] > 100:
+      value_color = 'PURPLE'
+      note = '(invalid?)'
+  else:
+    for threshold, color in ATTRIBUTE_COLORS:
+      threshold_val = dev.known_attributes[attr].get(threshold, None)
+      if threshold_val and value['raw'] >= threshold_val:
+        value_color = color
+        if threshold == 'Error':
+          note = '(failed)'
+        elif threshold == 'Maximum':
+          note = '(invalid?)'
+
+  # 199/C7 warning
+  if str(attr) == '199' and value['raw'] > 0:
+    note = '(bad cable?)'
+
+  # Done
+  return (value_color, note)
+
+
 def get_attribute_value_string(dev, attr) -> str:
   """Get attribute value string and report if it has changed."""
   current_value = dev.attributes.get(attr, {})