From 964885d63cf3d92b52af2ed321aba2c4a7a5295c Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 10 Nov 2019 17:38:03 -0700
Subject: [PATCH] Ensure tmux panes are closed atexit for hw-diags

---
 scripts/hw-diags.py    |  8 +-------
 scripts/wk/hw/diags.py | 10 ++++++++--
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/scripts/hw-diags.py b/scripts/hw-diags.py
index 6b34f324..e3719875 100755
--- a/scripts/hw-diags.py
+++ b/scripts/hw-diags.py
@@ -5,15 +5,9 @@
 import wk
 
 
-def main():
-  """Run hardware diagnostics."""
-  state = wk.hw.diags.State()
-  wk.hw.diags.main()
-
-
 if __name__ == '__main__':
   try:
-    main()
+    wk.hw.diags.main()
   except SystemExit:
     raise
   except: #pylint: disable=bare-except
diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py
index e61ee054..d1667c01 100644
--- a/scripts/wk/hw/diags.py
+++ b/scripts/wk/hw/diags.py
@@ -1,6 +1,7 @@
 """WizardKit: Hardware diagnostics"""
 # vim: sts=2 sw=2 ts=2
 
+import atexit
 import logging
 import pathlib
 import platform
@@ -14,6 +15,9 @@ from wk.cfg.hw import TMUX_SIDE_WIDTH
 from wk.cfg.main import KIT_NAME_FULL
 
 
+# atexit functions
+atexit.register(tmux.kill_all_panes)
+
 # STATIC VARIABLES
 DOCSTRING = f'''{KIT_NAME_FULL}: Hardware Diagnostics
 
@@ -102,7 +106,7 @@ class State():
       behind=True,
       lines=2,
       vertical=True,
-      text=self.top_text,
+      text=f'{self.top_text}\nMain Menu',
       )
 
     # Started
@@ -225,7 +229,6 @@ def main():
 
   # Show menu
   while True:
-    state.update_top_pane('Main Menu')
     action = None
     selection = menu.advanced_select()
 
@@ -257,6 +260,9 @@ def main():
       #run_diags()
       pass
 
+    # Reset top pane
+    state.update_top_pane('Main Menu')
+
 
 def network_test():
   """Run network tests."""