WizardKit/scripts/wk/hw/cpu.py

"""WizardKit: CPU test functions"""
# vim: sts=2 sw=2 ts=2

import logging
import re
import subprocess

from typing import TextIO

from wk import exe
from wk.cfg.hw import (
  CPU_TEMP_COOLING_DELTA,
  CPU_TEMP_CRITICAL,
  CPU_TEMP_IDLE_HIGH,
  CPU_TEMP_LOW_THRESHOLD,
  )
from wk.os.mac import set_fans as macos_set_fans
from wk.std import PLATFORM
from wk.ui import ansi


# STATIC VARIABLES
LOG = logging.getLogger(__name__)
SysbenchType = tuple[subprocess.Popen, TextIO]


# Functions
def check_cooling_results(sensors, test_object) -> None:
  """Check cooling result via sensor data."""
  idle_temp = sensors.get_cpu_temp('Idle')
  cooldown_temp = sensors.get_cpu_temp('Cooldown')
  max_temp = sensors.get_cpu_temp('Max')
  test_object.report.append(ansi.color_string('Temps', 'BLUE'))

  # Check temps
  if max_temp > CPU_TEMP_CRITICAL:
    test_object.failed = True
    test_object.set_status('Failed')
    test_object.report.extend([
      ansi.color_string(
        f'  WARNING: Critical CPU temp of {CPU_TEMP_CRITICAL} exceeded.',
        'RED',
        ),
      '',
      ])
  elif idle_temp >= CPU_TEMP_IDLE_HIGH:
    test_object.failed = True
    test_object.set_status('Failed')
    test_object.report.extend([
      ansi.color_string(
        f'  WARNING: Max idle temp of {CPU_TEMP_IDLE_HIGH} exceeded.',
        'YELLOW',
        ),
      '',
      ])
  elif (
      cooldown_temp <= CPU_TEMP_LOW_THRESHOLD
      or abs(max_temp - cooldown_temp) >= CPU_TEMP_COOLING_DELTA
      ):
    test_object.passed = True
    test_object.set_status('Passed')
  else:
    test_object.passed = False
    test_object.set_status('Unknown')

  # Build report
  report_labels = ['Idle']
  if 'Sysbench' in sensors.temp_labels:
    report_labels.extend(['Sysbench', 'Cooldown'])
  if 'Prime95' in sensors.temp_labels:
    report_labels.append('Prime95')
    if 'Cooldown' not in report_labels:
      report_labels.append('Cooldown')
  if len(sensors.temp_labels.intersection(['Prime95', 'Sysbench'])) < 1:
    # Include overall max temp if needed
    report_labels.append('Max')
  for line in sensors.generate_report(*report_labels, only_cpu=True):
    test_object.report.append(f'  {line}')


def check_mprime_results(test_obj, working_dir) -> None:
  """Check mprime log files and update test_obj."""
  passing_lines = {}
  warning_lines = {}

  def _read_file(log_name) -> list[str]:
    """Read file and split into lines, returns list."""
    lines = []
    try:
      with open(f'{working_dir}/{log_name}', 'r', encoding='utf-8') as _f:
        lines = _f.readlines()
    except FileNotFoundError:
      # File may be missing on older systems
      lines = []

    return lines

  # results.txt (check if failed)
  for line in _read_file('results.txt'):
    line = line.strip()
    if re.search(r'(error|fail)', line, re.IGNORECASE):
      warning_lines[line] = None

  # prime.log (check if passed)
  for line in _read_file('prime.log'):
    line = line.strip()
    match = re.search(
      r'(completed.*(\d+) errors, (\d+) warnings)', line, re.IGNORECASE)
    if match:
      if int(match.group(2)) + int(match.group(3)) > 0:
        # Errors and/or warnings encountered
        warning_lines[match.group(1).capitalize()] = None
      else:
        # No errors/warnings
        passing_lines[match.group(1).capitalize()] = None

  # Update status
  if warning_lines:
    test_obj.failed = True
    test_obj.set_status('Failed')
  elif passing_lines and 'Aborted' not in test_obj.status:
    test_obj.passed = True
    test_obj.set_status('Passed')
  else:
    test_obj.set_status('Unknown')

  # Update report
  for line in passing_lines:
    test_obj.report.append(f'  {line}')
  for line in warning_lines:
    test_obj.report.append(ansi.color_string(f'  {line}', 'YELLOW'))
  if not (passing_lines or warning_lines):
    test_obj.report.append(ansi.color_string('  Unknown result', 'YELLOW'))


def start_mprime(working_dir, log_path) -> subprocess.Popen:
  """Start mprime and save filtered output to log, returns Popen object."""
  set_apple_fan_speed('max')
  proc_mprime = subprocess.Popen(
    ['mprime', '-t'],
    cwd=working_dir,
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    )
  proc_grep = subprocess.Popen(
    'grep --ignore-case --invert-match --line-buffered stress.txt'.split(),
    stdin=proc_mprime.stdout,
    stdout=subprocess.PIPE,
    )
  proc_mprime.stdout.close() # type: ignore[reportOptionalMemberAccess]
  save_nbsr = exe.NonBlockingStreamReader(
    proc_grep.stdout, # type: ignore[reportGeneralTypeIssues]
    )
  exe.start_thread(
    save_nbsr.save_to_file,
    args=(proc_grep, log_path),
    )

  # Return objects
  return proc_mprime


def set_apple_fan_speed(speed) -> None:
  """Set Apple fan speed."""
  cmd = None

  # Check
  if speed not in ('auto', 'max'):
    raise RuntimeError(f'Invalid speed {speed}')

  # Set cmd
  if PLATFORM == 'Darwin':
    try:
      macos_set_fans(speed)
    except (RuntimeError, ValueError, subprocess.CalledProcessError) as err:
      LOG.error('Failed to set fans to %s', speed)
      LOG.error('Error: %s', err)
      #ui.print_error(f'Failed to set fans to {speed}')
      #for line in str(err).splitlines():
      #  ui.print_warning(f'  {line.strip()}')
  elif PLATFORM == 'Linux':
    cmd = ['apple-fans', speed]
    exe.run_program(cmd, check=False)


def start_sysbench(log_path) -> SysbenchType:
  """Start sysbench, returns tuple with Popen object and file handle."""
  set_apple_fan_speed('max')
  cmd = [
    'sysbench',
    f'--threads={exe.psutil.cpu_count()}',
    '--cpu-max-prime=1000000000',
    'cpu',
    'run',
    ]

  # Start sysbench
  filehandle = open(
    log_path, 'a', encoding='utf-8',
    )
  proc = exe.popen_program(cmd, stdout=filehandle)

  # Done
  return (proc, filehandle)


def stop_mprime(proc_mprime) -> None:
  """Stop mprime gracefully, then forcefully as needed."""
  proc_mprime.terminate()
  try:
    proc_mprime.wait(timeout=5)
  except subprocess.TimeoutExpired:
    proc_mprime.kill()
  set_apple_fan_speed('auto')


def stop_sysbench(proc_sysbench, filehandle_sysbench) -> None:
  """Stop sysbench."""
  proc_sysbench.terminate()
  try:
    proc_sysbench.wait(timeout=5)
  except subprocess.TimeoutExpired:
    proc_sysbench.kill()
  filehandle_sysbench.flush()
  filehandle_sysbench.close()
  set_apple_fan_speed('auto')


if __name__ == '__main__':
  print("This file is not meant to be called directly.")