Track test groups as list instead of a dict

By only including selected tests we can skip checking for the
enabled/disabled status.  This will also simplify the process
of disabling future tests for a disk if a failure is detected.

scripts/wk/cfg/hw.py

@@ -16,9 +16,21 @@ ATTRIBUTE_COLORS = (
   )
 # NOTE: Force 4K read block size for disks >= 3TB
 BADBLOCKS_LARGE_DISK = 3 * 1024**4
+BADBLOCKS_REGEX = re.compile(
+  r'^Pass completed, (\d+) bad blocks found. .(\d+)/(\d+)/(\d+) errors',
+  re.IGNORECASE,
+  )
 CPU_CRITICAL_TEMP = 99
 CPU_FAILURE_TEMP = 90
 CPU_TEST_MINUTES = 7
+IO_GRAPH_WIDTH = 40
+IO_ALT_TEST_SIZE_FACTOR = 0.01
+IO_BLOCK_SIZE = 512 * 1024
+IO_CHUNK_SIZE = 32 * 1024**2
+IO_MINIMUM_TEST_SIZE = 10 * 1024**3
+IO_RATE_REGEX = re.compile(
+  r'(?P<bytes>\d+) bytes.* (?P<seconds>\S+) s(?:,|ecs )',
+  )
 KEY_NVME = 'nvme_smart_health_information_log'
 KEY_SMART = 'ata_smart_attributes'
 KNOWN_DISK_ATTRIBUTES = {
@@ -116,6 +128,18 @@ SMC_IDS = {
   'Tp5P': {'CPU Temp': False, 'Source': 'PSU2 Secondary Component'},
   'TS0C': {'CPU Temp': False, 'Source': 'CPU B DIMM Exit Ambient'},
   }
+STATUS_COLORS = {
+  'Passed': 'GREEN',
+  'Aborted': 'YELLOW',
+  'N/A': 'YELLOW',
+  'Skipped': 'YELLOW',
+  'Unknown': 'YELLOW',
+  'Working': 'YELLOW',
+  'Denied': 'RED',
+  'ERROR': 'RED',
+  'Failed': 'RED',
+  'TimedOut': 'RED',
+  }
 TEMP_COLORS = {
   float('-inf'):  'CYAN',
   00:             'BLUE',

scripts/wk/hw/diags.py

@@ -10,15 +10,24 @@ import re
 import subprocess
 import time
 
-from collections import OrderedDict
 from docopt import docopt
 
 from wk import cfg, debug, exe, graph, log, net, std, tmux
 from wk import os as wk_os
+from wk.cfg.hw import (
+  BADBLOCKS_REGEX,
+  IO_GRAPH_WIDTH,
+  IO_ALT_TEST_SIZE_FACTOR,
+  IO_BLOCK_SIZE,
+  IO_CHUNK_SIZE,
+  IO_MINIMUM_TEST_SIZE,
+  IO_RATE_REGEX,
+  STATUS_COLORS,
+  )
 from wk.hw import disk as hw_disk
 from wk.hw import sensors as hw_sensors
 from wk.hw import system as hw_system
-from wk.hw.test import Test
+from wk.hw.test import Test, TestGroup
 
 
 # STATIC VARIABLES
@@ -34,18 +43,15 @@ Options:
   -q --quick          Skip menu and perform a quick check
 '''
 LOG = logging.getLogger(__name__)
-BADBLOCKS_REGEX = re.compile(
+TEST_GROUPS = {
-  r'^Pass completed, (\d+) bad blocks found. .(\d+)/(\d+)/(\d+) errors',
+  # Also used to build the menu options
-  re.IGNORECASE,
+  ## NOTE: This needs to be above MENU_SETS
-  )
+  'CPU & Cooling': 'cpu_stress_tests',
-IO_GRAPH_WIDTH = 40
+  'Disk Attributes': 'disk_attribute_check',
-IO_ALT_TEST_SIZE_FACTOR = 0.01
+  'Disk Self-Test': 'disk_self_test',
-IO_BLOCK_SIZE = 512 * 1024
+  'Disk Surface Scan': 'disk_surface_scan',
-IO_CHUNK_SIZE = 32 * 1024**2
+  'Disk I/O Benchmark': 'disk_io_benchmark',
-IO_MINIMUM_TEST_SIZE = 10 * 1024**3
+  }
-IO_RATE_REGEX = re.compile(
-  r'(?P<bytes>\d+) bytes.* (?P<seconds>\S+) s(?:,|ecs )',
-  )
 MENU_ACTIONS = (
   'Audio Test',
   'Keyboard Test',
@@ -57,16 +63,9 @@ MENU_ACTIONS_SECRET = (
   'Matrix',
   'Tubes',
   )
-MENU_OPTIONS = (
-  'CPU & Cooling',
-  'Disk Attributes',
-  'Disk Self-Test',
-  'Disk Surface Scan',
-  'Disk I/O Benchmark',
-)
 MENU_OPTIONS_QUICK = ('Disk Attributes',)
 MENU_SETS = {
-  'Full Diagnostic': (*MENU_OPTIONS,),
+  'Full Diagnostic': (*TEST_GROUPS,),
   'Disk Diagnostic': (
     'Disk Attributes',
     'Disk Self-Test',
@@ -79,19 +78,6 @@ MENU_TOGGLES = (
   'Skip USB Benchmarks',
   )
 PLATFORM = std.PLATFORM
-STATUS_COLORS = {
-  'Passed': 'GREEN',
-  'Aborted': 'YELLOW',
-  'N/A': 'YELLOW',
-  'Skipped': 'YELLOW',
-  'Unknown': 'YELLOW',
-  'Working': 'YELLOW',
-  'Denied': 'RED',
-  'ERROR': 'RED',
-  'Failed': 'RED',
-  'TimedOut': 'RED',
-  }
-
 
 # Error Classes
 class DeviceTooSmallError(RuntimeError):
@@ -107,33 +93,7 @@ class State():
     self.log_dir = None
     self.panes = {}
     self.system = None
-    self.tests = OrderedDict({
+    self.test_groups = []
-      'CPU & Cooling': {
-        'Enabled': False,
-        'Function': cpu_stress_tests,
-        'Objects': [],
-        },
-      'Disk Attributes': {
-        'Enabled': False,
-        'Function': disk_attribute_check,
-        'Objects': [],
-        },
-      'Disk Self-Test': {
-        'Enabled': False,
-        'Function': disk_self_test,
-        'Objects': [],
-        },
-      'Disk Surface Scan': {
-        'Enabled': False,
-        'Function': disk_surface_scan,
-        'Objects': [],
-        },
-      'Disk I/O Benchmark': {
-        'Enabled': False,
-        'Function': disk_io_benchmark,
-        'Objects': [],
-        },
-      })
     self.top_text = std.color_string('Hardware Diagnostics', 'GREEN')
 
     # Init tmux and start a background process to maintain layout
@@ -142,8 +102,8 @@ class State():
 
   def abort_testing(self):
     """Set unfinished tests as aborted and cleanup tmux panes."""
-    for details in self.tests.values():
+    for group in self.test_groups:
-      for test in details['Objects']:
+      for test in group.test_objects:
         if test.status in ('Pending', 'Working'):
           test.set_status('Aborted')
 
@@ -269,8 +229,7 @@ class State():
     self.disks.clear()
     self.layout.clear()
     self.layout.update(cfg.hw.TMUX_LAYOUT)
-    for test_data in self.tests.values():
+    self.test_groups.clear()
-      test_data['Objects'].clear()
 
     # Set log
     self.log_dir = log.format_log_path()
@@ -300,9 +259,10 @@ class State():
 
     # Add test objects
     for name, details in menu.options.items():
-      self.tests[name]['Enabled'] = details['Selected']
       if not details['Selected']:
+        # Only add selected options
         continue
+
       if 'CPU' in name:
         # Create two Test objects which will both be used by cpu_stress_tests
         # NOTE: Prime95 should be added first
@@ -312,12 +272,23 @@ class State():
         self.system.tests.append(
           Test(dev=self.system, label='Cooling', name=name),
           )
-        self.tests[name]['Objects'].extend(self.system.tests)
+        self.test_groups.append(
-      elif 'Disk' in name:
+          TestGroup(
+            name=name,
+            function=globals()[TEST_GROUPS[name]],
+            test_objects=self.system.tests,
+            ),
+          )
+
+      if 'Disk' in name:
+        test_group = TestGroup(
+          name=name, function=globals()[TEST_GROUPS[name]],
+          )
         for disk in self.disks:
           test_obj = Test(dev=disk, label=disk.path.name, name=name)
           disk.tests.append(test_obj)
-          self.tests[name]['Objects'].append(test_obj)
+          test_group.test_objects.append(test_obj)
+        self.test_groups.append(test_group)
 
     # Run safety checks
     self.disk_safety_checks(prep=True)
@@ -413,16 +384,12 @@ class State():
     report = []
     width = cfg.hw.TMUX_SIDE_WIDTH
 
-    for name, details in self.tests.items():
+    for group in self.test_groups:
-      if not details['Enabled']:
+      report.append(std.color_string(group.name, 'BLUE'))
-        continue
+      for test in group.test_objects:
-
-      # Add test details
-      report.append(std.color_string(name, 'BLUE'))
-      for test_obj in details['Objects']:
         report.append(std.color_string(
-          [test_obj.label, f'{test_obj.status:>{width-len(test_obj.label)}}'],
+          [test.label, f'{test.status:>{width-len(test.label)}}'],
-          [None, STATUS_COLORS.get(test_obj.status, None)],
+          [None, STATUS_COLORS.get(test.status, None)],
           sep='',
           ))
 
@@ -471,7 +438,7 @@ def build_menu(cli_mode=False, quick_mode=False):
     menu.add_action(action)
   for action in MENU_ACTIONS_SECRET:
     menu.add_action(action, {'Hidden': True})
-  for option in MENU_OPTIONS:
+  for option in TEST_GROUPS:
     menu.add_option(option, {'Selected': True})
   for toggle in MENU_TOGGLES:
     menu.add_toggle(toggle, {'Selected': True})
@@ -1362,21 +1329,18 @@ def run_diags(state, menu, quick_mode=False):
   state.init_diags(menu)
 
   # Just return if no tests were selected
-  if not any(details['Enabled'] for details in state.tests.values()):
+  if not state.test_groups:
     std.print_warning('No tests selected?')
     std.pause()
     return
 
   # Run tests
-  for name, details in state.tests.items():
+  for group in state.test_groups:
-    if not details['Enabled']:
-      # Skip disabled tests
-      continue
 
     # Run test(s)
-    function = details['Function']
+    function = group.function
-    args = [details['Objects']]
+    args = [group.test_objects]
-    if name == 'Disk I/O Benchmark':
+    if group.name == 'Disk I/O Benchmark':
       args.append(menu.toggles['Skip USB Benchmarks']['Selected'])
     std.clear_screen()
     try:
@@ -1388,15 +1352,17 @@ def run_diags(state, menu, quick_mode=False):
       break
 
     # Run safety checks
-    if name.startswith('Disk'):
+    if group.name.startswith('Disk'):
-      state.disk_safety_checks(wait_for_self_tests=name != 'Disk Attributes')
+      state.disk_safety_checks(
+        wait_for_self_tests=group.name != 'Disk Attributes',
+        )
 
   # Handle aborts
   if aborted:
-    for details in state.tests.values():
+    for group in state.test_groups:
-      for test_obj in details['Objects']:
+      for test in group.test_objects:
-        if test_obj.status == 'Pending':
+        if test.status == 'Pending':
-          test_obj.set_status('Aborted')
+          test.set_status('Aborted')
 
   # Show results
   show_results(state)
@@ -1462,17 +1428,19 @@ def show_results(state):
   state.update_top_pane('Results')
 
   # CPU Tests
-  cpu_tests_enabled = [data['Enabled'] for name, data in state.tests.items()
+  cpu_tests_enabled = [
-                       if name.startswith('CPU')]
+    group.name for group in state.test_groups if 'CPU' in group.name
-  if any(cpu_tests_enabled):
+    ]
+  if cpu_tests_enabled:
     std.print_success('CPU:')
     std.print_report(state.system.generate_report())
     std.print_standard(' ')
 
   # Disk Tests
-  disk_tests_enabled = [data['Enabled'] for name, data in state.tests.items()
+  disk_tests_enabled = [
-                        if name.startswith('Disk')]
+    group.name for group in state.test_groups if 'Disk' in group.name
-  if any(disk_tests_enabled):
+    ]
+  if disk_tests_enabled:
     std.print_success(f'Disk{"s" if len(state.disks) > 1 else ""}:')
     for disk in state.disks:
       std.print_report(disk.generate_report())

scripts/wk/hw/test.py

@@ -2,7 +2,7 @@
 # vim: sts=2 sw=2 ts=2
 
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Any, Callable
 
 @dataclass(slots=True)
 class Test:
@@ -25,3 +25,11 @@ class Test:
       return
 
     self.status = status
+
+
+@dataclass(slots=True)
+class TestGroup:
+  """Object for tracking groups of tests."""
+  name: str
+  function: Callable
+  test_objects: list[Test] = field(default_factory=list)