From 312df1ee9dac468d7483855c1931a9edf615fe32 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Fri, 8 Apr 2022 15:38:03 -0600
Subject: [PATCH] Refactor import logic for wk.hw.diags

---
 scripts/wk/hw/diags.py | 29 ++++++++++-------------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py
index b4fccb13..55d9f53c 100644
--- a/scripts/wk/hw/diags.py
+++ b/scripts/wk/hw/diags.py
@@ -16,21 +16,12 @@ from wk.hw import benchmark as hw_benchmark
 from wk.hw import cpu as hw_cpu
 from wk.hw import disk as hw_disk
 from wk.hw import sensors as hw_sensors
+from wk.hw import smart as hw_smart
 from wk.hw import surface_scan as hw_surface_scan
 from wk.hw import system as hw_system
 from wk.hw.audio import audio_test
 from wk.hw.keyboard import keyboard_test
 from wk.hw.network import network_test
-from wk.hw.smart import (
-  CriticalHardwareError,
-  SMARTSelfTestInProgressError,
-  abort_self_test,
-  check_attributes,
-  check_self_test_results,
-  generate_attribute_report,
-  run_self_test,
-  safety_checks,
-  )
 from wk.hw.screensavers import screensaver
 from wk.hw.test import Test, TestGroup
 
@@ -131,8 +122,8 @@ class State():
         continue
 
       try:
-        safety_checks(disk)
-      except CriticalHardwareError:
+        hw_smart.safety_checks(disk)
+      except hw_smart.CriticalHardwareError:
         disable_tests = True
         disk.add_note('Critical hardware error detected.', 'RED')
         if 'Disk Attributes' in disk.tests:
@@ -145,7 +136,7 @@ class State():
             'Critical hardware error detected during diagnostics',
             'YELLOW',
             )
-      except SMARTSelfTestInProgressError as err:
+      except hw_smart.SMARTSelfTestInProgressError as err:
         if prep:
           std.print_warning(f'SMART self-test(s) in progress for {disk.path}')
           if std.ask('Continue with all tests disabled for this device?'):
@@ -170,7 +161,7 @@ class State():
         if (
             'Disk Attributes' in disk.tests
             and not disk.tests['Disk Attributes'].failed
-            and not check_attributes(disk, only_blocking=False)
+            and not hw_smart.check_attributes(disk, only_blocking=False)
             ):
           # No blocking errors encountered, but found minor attribute failures
           if not prep:
@@ -595,7 +586,7 @@ def disk_attribute_check(state, test_objects) -> None:
       test.set_status('N/A')
       continue
 
-    if check_attributes(test.dev):
+    if hw_smart.check_attributes(test.dev):
       test.passed = True
       test.set_status('Passed')
     else:
@@ -691,7 +682,7 @@ def disk_self_test(state, test_objects) -> None:
     # Start thread
     test.set_status('Working')
     test_log = f'{state.log_dir}/{test.dev.path.name}_selftest.log'
-    threads.append(exe.start_thread(run_self_test, args=(test, test_log)))
+    threads.append(exe.start_thread(hw_smart.run_self_test, args=(test, test_log)))
 
     # Show progress
     if threads[-1].is_alive():
@@ -710,12 +701,12 @@ def disk_self_test(state, test_objects) -> None:
   except KeyboardInterrupt:
     aborted = True
     for test in test_objects:
-      abort_self_test(test.dev)
+      hw_smart.abort_self_test(test.dev)
     std.sleep(0.5)
 
   # Save report(s)
   for test in test_objects:
-    check_self_test_results(test, aborted=aborted)
+    hw_smart.check_self_test_results(test, aborted=aborted)
 
   # Cleanup
   state.update_progress_pane()
@@ -745,7 +736,7 @@ def disk_surface_scan(state, test_objects) -> None:
     )
   for disk in state.disks:
     failed_attributes = [
-      line for line in generate_attribute_report(disk) if 'failed' in line
+      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)