From 60168e14aeaf50f7995329aa87b29bf51aa905c0 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 3 Dec 2022 20:50:09 -0800
Subject: [PATCH] Fix badblocks reports

- Better filter the raw results, stripping the backspaces from the text.
- Fix color used when printing to the screen
---
 scripts/wk/cfg/hw.py          |  4 +++-
 scripts/wk/hw/surface_scan.py | 30 +++++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/scripts/wk/cfg/hw.py b/scripts/wk/cfg/hw.py
index 45483ed6..f4090ae5 100644
--- a/scripts/wk/cfg/hw.py
+++ b/scripts/wk/cfg/hw.py
@@ -21,7 +21,9 @@ BADBLOCKS_REGEX = re.compile(
   r'^Pass completed, (\d+) bad blocks found. .(\d+)/(\d+)/(\d+) errors',
   re.IGNORECASE,
   )
-BADBLOCKS_RESULTS_REGEX = re.compile(r'^(.*?)\x08.*\x08(.*)')
+BADBLOCKS_RESULTS_REGEX = re.compile(
+  r'^(Checking for bad blocks .read-only test.: ).*\x08+(done|\s+).*?(\x08+)?'
+  )
 BADBLOCKS_SKIP_REGEX = re.compile(r'^(Checking|\[)', re.IGNORECASE)
 CPU_CRITICAL_TEMP = 99
 CPU_FAILURE_TEMP = 90
diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py
index 0665ac99..a7e10fb8 100644
--- a/scripts/wk/hw/surface_scan.py
+++ b/scripts/wk/hw/surface_scan.py
@@ -29,31 +29,47 @@ LOG = logging.getLogger(__name__)
 # Functions
 def check_surface_scan_results(test_obj, log_path) -> None:
   """Check results and set test status."""
+  report = []
+  report_color = None
+
+  # Read result
   with open(log_path, 'r', encoding='utf-8') as _f:
     for line in _f.readlines():
-      # TODO: Test further to restrict report to just the results
       line = strip_colors(line.strip())
-      #if not line or BADBLOCKS_SKIP_REGEX.match(line):
       if not line:
         # Skip
         continue
 
       # Clean line by removing backspaces/etc
-      line = BADBLOCKS_RESULTS_REGEX.sub(r'\1 \2', line)
+      match = BADBLOCKS_RESULTS_REGEX.match(line)
+      if match:
+        if match.group(2) == ' done':
+          line = BADBLOCKS_RESULTS_REGEX.sub(r'\1done', line)
+        else:
+          line = BADBLOCKS_RESULTS_REGEX.sub(r'\1\2', line)
+      line = line.replace(r'\x08', '')
+      line = line.strip()
 
       # Add to report
+      report.append(line)
       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")}')
+
+  # Set report color and save to test_obj
+  if test_obj.failed:
+    report_color = 'RED'
+  elif not test_obj.passed:
+    report_color = 'YELLOW'
+  for line in report:
+    test_obj.report.append(f'  {color_string(line, report_color)}')
+
+  # Handle undefined result status
   if not (test_obj.passed or test_obj.failed):
     test_obj.set_status('Unknown')