From d667695e9eafa145f14c8bb4b64fbe0e208e95d1 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Mon, 11 Apr 2022 14:30:42 -0600
Subject: [PATCH] Move surface scan check to its own function

---
 scripts/wk/cfg/hw.py          |  1 +
 scripts/wk/hw/surface_scan.py | 52 ++++++++++++++++++++---------------
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/scripts/wk/cfg/hw.py b/scripts/wk/cfg/hw.py
index ae39da69..e0ae4cd8 100644
--- a/scripts/wk/cfg/hw.py
+++ b/scripts/wk/cfg/hw.py
@@ -20,6 +20,7 @@ BADBLOCKS_REGEX = re.compile(
   r'^Pass completed, (\d+) bad blocks found. .(\d+)/(\d+)/(\d+) errors',
   re.IGNORECASE,
   )
+BADBLOCKS_SKIP_REGEX = re.compile(r'^(Checking|\[)', re.IGNORECASE)
 CPU_CRITICAL_TEMP = 99
 CPU_FAILURE_TEMP = 90
 CPU_TEST_MINUTES = 7
diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py
index 85d87f96..f3a7212d 100644
--- a/scripts/wk/hw/surface_scan.py
+++ b/scripts/wk/hw/surface_scan.py
@@ -5,7 +5,11 @@ import logging
 
 from subprocess import STDOUT
 
-from wk.cfg.hw import BADBLOCKS_LARGE_DISK, BADBLOCKS_REGEX
+from wk.cfg.hw import (
+  BADBLOCKS_LARGE_DISK,
+  BADBLOCKS_REGEX,
+  BADBLOCKS_SKIP_REGEX,
+  )
 from wk.exe import run_program
 from wk.std import (
   PLATFORM,
@@ -20,6 +24,30 @@ LOG = logging.getLogger(__name__)
 
 
 # Functions
+def check_surface_scan_results(test_obj, log_path) -> None:
+  """Check results and set test status."""
+  with open(log_path, 'r', encoding='utf-8') as _f:
+    for line in _f.readlines():
+      line = strip_colors(line.strip())
+      if not line or BADBLOCKS_SKIP_REGEX.match(line):
+        # Skip
+        continue
+      match = BADBLOCKS_REGEX.search(line)
+      if match:
+        if all(s == '0' for s in match.groups()):
+          test_obj.passed = True
+          test_obj.report.append(f'  {line}')
+          test_obj.set_status('Passed')
+        else:
+          test_obj.failed = True
+          test_obj.report.append(f'  {color_string(line, "YELLOW")}')
+          test_obj.set_status('Failed')
+      else:
+        test_obj.report.append(f'  {color_string(line, "YELLOW")}')
+  if not (test_obj.passed or test_obj.failed):
+    test_obj.set_status('Unknown')
+
+
 def run_scan(test_obj, log_path) -> None:
   """Run surface scan and handle exceptions."""
   block_size = '1024'
@@ -58,27 +86,7 @@ def run_scan(test_obj, log_path) -> None:
       )
 
   # Check results
-  with open(log_path, 'r', encoding='utf-8') as _f:
-    for line in _f.readlines():
-      line = strip_colors(line.strip())
-      if not line or line.startswith('Checking') or line.startswith('['):
-        # Skip
-        continue
-      match = BADBLOCKS_REGEX.search(line)
-      if match:
-        if all(s == '0' for s in match.groups()):
-          test_obj.passed = True
-          test_obj.report.append(f'  {line}')
-          test_obj.set_status('Passed')
-        else:
-          test_obj.failed = True
-          test_obj.report.append(f'  {color_string(line, "YELLOW")}')
-          test_obj.set_status('Failed')
-      else:
-        test_obj.report.append(f'  {color_string(line, "YELLOW")}')
-  if not (test_obj.passed or test_obj.failed):
-    test_obj.set_status('Unknown')
-
+  check_surface_scan_results(test_obj, log_path)
 
 
 if __name__ == '__main__':