From a28a754be1621b63e008a31bc0e7348b8fc0bdb1 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Tue, 5 Feb 2019 21:08:06 -0700
Subject: [PATCH] Added thermal limit logic to Prime95 test

---
 .bin/Scripts/functions/hw_diags.py | 24 +++++++++++++++++++-----
 .bin/Scripts/settings/main.py      |  1 +
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/.bin/Scripts/functions/hw_diags.py b/.bin/Scripts/functions/hw_diags.py
index b1ed5869..44a1c7f7 100644
--- a/.bin/Scripts/functions/hw_diags.py
+++ b/.bin/Scripts/functions/hw_diags.py
@@ -1262,6 +1262,7 @@ def run_mprime_test(state, test):
   test.update_status()
   update_progress_pane(state)
   test.sensor_data = get_sensor_data()
+  test.thermal_abort = False
 
   # Update tmux layout
   tmux_update_pane(
@@ -1303,6 +1304,7 @@ def run_mprime_test(state, test):
     command=['hw-diags-prime95', global_vars['TmpDir']],
     working_dir=global_vars['TmpDir'])
   time_limit = int(MPRIME_LIMIT) * 60
+  thermal_limit = int(THERMAL_LIMIT)
   try:
     for i in range(time_limit):
       clear_screen()
@@ -1319,15 +1321,19 @@ def run_mprime_test(state, test):
       # Not using print wrappers to avoid flooding the log
       print(_status_str)
       print('{YELLOW}{msg}{CLEAR}'.format(msg=test.abort_msg, **COLORS))
-      update_sensor_data(test.sensor_data)
+      update_sensor_data(test.sensor_data, thermal_limit)
 
       # Wait
       sleep(1)
-  except KeyboardInterrupt:
-    # Catch CTRL+C
+  except (KeyboardInterrupt, ThermalError) as err:
+    # CTRL+c pressed or thermal threshold reached
     test.aborted = True
-    test.update_status('Aborted')
-    print_warning('\nAborted.')
+    if isinstance(err, KeyboardInterrupt):
+      test.update_status('Aborted')
+    elif isinstance(err, ThermalError):
+      test.failed = True
+      test.thermal_abort = True
+      test.update_status('NS')
     update_progress_pane(state)
 
     # Restart live monitor
@@ -1428,6 +1434,14 @@ def run_mprime_test(state, test):
       test.sensor_data, 'Idle', 'Max', 'Cooldown', core_only=True):
     test.report.append('  {}'.format(line))
 
+  # Add abort message(s)
+  if test.aborted:
+    test.report.append(
+      '  {YELLOW}Aborted{CLEAR}'.format(**COLORS))
+  if test.thermal_abort:
+    test.report.append(
+      '  {RED}CPU temps exceeded threshold{CLEAR}'.format(**COLORS))
+
   # Done
   update_progress_pane(state)
 
diff --git a/.bin/Scripts/settings/main.py b/.bin/Scripts/settings/main.py
index da83e830..e824f095 100644
--- a/.bin/Scripts/settings/main.py
+++ b/.bin/Scripts/settings/main.py
@@ -16,6 +16,7 @@ KIT_NAME_SHORT='WK'
 SUPPORT_MESSAGE='Please let 2Shirt know by opening an issue on GitHub'
 # Live Linux
 MPRIME_LIMIT='7'                            # of minutes to run Prime95 during hw-diags
+THERMAL_LIMIT='95'                          # Prime95 abort temperature in Celsius
 ROOT_PASSWORD='Abracadabra'
 TECH_PASSWORD='Abracadabra'
 # Server IP addresses