WizardKit/.bin/Scripts/hw-diags-menu

#!/bin/python3
#
## Wizard Kit: HW Diagnostics - Menu

import libtmux
import json
import os
import sys

# Init
os.chdir(os.path.dirname(os.path.realpath(__file__)))
sys.path.append(os.getcwd())
from functions.common import *
init_global_vars()

# STATIC VARIABLES
sleep(1)
TMUX = libtmux.Server()
SESSION = TMUX.find_where({'session_name': 'hw-diags'})
WINDOW = SESSION.windows[0] # Should be a safe assumption
PANE = WINDOW.panes[0]      # Should be a safe assumption
PROGRESS_FILE = '{}/progress.out'.format(global_vars['LogDir'])
TESTS = {
    'Prime95': {
        'Enabled': False,
        'Status':  'Pending',
        },
    'SMART': {
        'Enabled': False,
        'Quick':   False,
        },
    'badblocks': {
        'Enabled': False,
        },
    }

def get_status_color(s):
    color = COLORS['CLEAR']
    if s in ['NS', 'Unknown']:
        color = COLORS['RED']
    elif s in ['Working', 'Skipped']:
        color = COLORS['YELLOW']
    elif s in ['CS']:
        color = COLORS['GREEN']
    return color

def menu_diags():
    diag_modes = [
        {'Name': 'All tests', 'Tests': ['Prime95', 'SMART', 'badblocks']},
        {'Name': 'Prime95', 'Tests': ['Prime95']},
        {'Name': 'SMART & badblocks', 'Tests': ['SMART', 'badblocks']},
        {'Name': 'SMART', 'Tests': ['SMART']},
        {'Name': 'badblocks', 'Tests': ['badblocks']},
        {'Name': 'Quick drive test', 'Tests': ['Quick', 'SMART']},
        ]
    actions = [
        {'Letter': 'A', 'Name': 'Audio test'},
        {'Letter': 'N', 'Name': 'Network test'},
        {'Letter': 'M', 'Name': 'Screen Saver - Matrix', 'CRLF': True},
        {'Letter': 'P', 'Name': 'Screen Saver - Pipes'},
        {'Letter': 'Q', 'Name': 'Quit', 'CRLF': True},
        ]

    # Show menu
    while True:
        selection = menu_select(
            title = 'Hardware Diagnostics: Menu',
            main_entries = diag_modes,
            action_entries = actions,
            spacer = '─────────────────────────')
        if selection.isnumeric():
            run_tests(diag_modes[int(selection)-1]['Tests'])
        elif selection == 'A':
            run_program(['hw-diags-audio'], check=False, pipe=False)
            sleep(1)
        elif selection == 'N':
            run_program(['hw-diags-network'], check=False, pipe=False)
            sleep(1)
        elif selection == 'M':
            run_program(['cmatrix', '-abs'], check=False, pipe=False)
        elif selection == 'P':
            run_program(
                'pipes -t 0 -t 1 -t 2 -t 3 -p 5 -R -r 4000'.split(),
                check=False, pipe=False)
        elif selection == 'Q':
            break

def run_badblocks():
    pass

def run_mprime():
    # Set Window layout
    window = SESSION.new_window()
    pane_sensors = window.panes[0]
    pane_mprime = window.split_window(attach=False)
    pane_mprime.set_height(10)
    pane_progress = window.split_window(attach=False, vertical=False)
    pane_progress.set_width(16)
    
    # Start test
    run_program(['apple-fans', 'max'])
    pane_sensors.send_keys('watch -c -n1 -t hw-sensors')
    pane_progress.send_keys('watch -c -n1 -t cat "{}"'.format(PROGRESS_FILE))
    pane_mprime.send_keys('mprime -t')
    #sleep(MPRIME_LIMIT*60)
    sleep(15)

    # Done
    run_program(['apple-fans', 'auto'])
    window.kill_window()

def run_smart():
    pass

def run_tests(tests):
    for t in ['Prime95', 'SMART', 'badblocks']:
        TESTS[t]['Enabled'] = t in tests
    TESTS['SMART']['Quick'] = 'Quick' in tests

    # Get (disk) device list
    if TESTS['SMART']['Enabled'] or TESTS['badblocks']['Enabled']:
        cmd = 'lsblk -J -o NAME,TYPE'.split()
        result = run_program(cmd)
        json_data = json.loads(result.stdout.decode())
        devs = json_data.get('blockdevices', [])
        devs = {d['name']: 'Pending' for d in devs if d['type'] == 'disk'}
        TESTS['SMART']['Devices'] = devs
        TESTS['badblocks']['Devices'] = devs

    # Initialize progress display
    update_progress()

    # Run
    if TESTS['Prime95']['Enabled']:
        run_mprime()
    if TESTS['SMART']['Enabled']:
        run_smart()
    if TESTS['badblocks']['Enabled']:
        run_badblocks()

def update_progress(color=True):
    output = []
    if color:
        output.append('{BLUE}HW  Diagnostics{CLEAR}'.format(**COLORS))
        output.append('───────────────')
        if TESTS['Prime95']['Enabled']:
            output.append('{BLUE}Prime95{s_color}{status:>8}{CLEAR}'.format(
                s_color = get_status_color(TESTS['Prime95']['Status']),
                status = TESTS['Prime95']['Status'],
                **COLORS))
        if TESTS['SMART']['Enabled']:
            output.append('{BLUE}SMART{CLEAR}'.format(**COLORS))
            for dev, status in sorted(TESTS['SMART']['Devices'].items()):
                output.append('{dev}{s_color}{status:>{pad}}{CLEAR}'.format(
                    dev = dev,
                    pad = 16-len(dev),
                    s_color = get_status_color(status),
                    status = status,
                    **COLORS))
        if TESTS['badblocks']['Enabled']:
            output.append('{BLUE}badblocks{CLEAR}'.format(**COLORS))
            for dev, status in sorted(TESTS['badblocks']['Devices'].items()):
                output.append('{dev}{s_color}{status:>{pad}}{CLEAR}'.format(
                    dev = dev,
                    pad = 16-len(dev),
                    s_color = get_status_color(status),
                    status = status,
                    **COLORS))
    else:
        output.append('HW  Diagnostics')
        output.append('───────────────')
        if TESTS['Prime95']['Enabled']:
            output.append('Prime95{:>8}'.format(TESTS['Prime95']['Status']))
        if TESTS['SMART']['Enabled']:
            output.append('SMART')
            for dev, status in sorted(TESTS['SMART']['Devices'].items()):
                output.append('{}{:>{}}'.format(dev, 16-len(dev), status))
        if TESTS['badblocks']['Enabled']:
            output.append('badblocks')
            for dev, status in sorted(TESTS['badblocks']['Devices'].items()):
                output.append('{}{:>{}}'.format(dev, 16-len(dev), status))

    # Add line-endings
    output = ['{}\n'.format(line) for line in output]
    print(output)
    pause()

    with open(PROGRESS_FILE, 'w') as f:
        f.writelines(output)

if __name__ == '__main__':
    try:
        # Prep
        clear_screen()

        # Show menu
        diags_menu()
        
        # Done
        #print_standard('\nDone.')
        #pause("Press Enter to exit...")
        exit_script()
    except SystemExit:
        pass
    except:
        major_exception()