crosperf: migrate all device setup code to a separate utils file

This patch extracts all the device setup code which interacts with DUT
to a single utils file, and be put into a wrapper class. This will help
migrating all related code to telemetry_Crosperf for skylab runs.

BUG=chromium:1020655
TEST=Passed all unittests; tested with simple experiment on kevin.

Change-Id: I2edcd7bb2d8cd0255d3ae6d380a5983c24427d98
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1895500
Tested-by: Zhizhou Yang <zhizhouy@google.com>
Reviewed-by: Zhizhou Yang <zhizhouy@google.com>
Reviewed-by: Denis Nikitin <denik@chromium.org>
Commit-Queue: Zhizhou Yang <zhizhouy@google.com>

2 files changed, 1073 insertions, 0 deletions

diff --git a/cros_utils/device_setup_utils.py b/cros_utils/device_setup_utils.py
new file mode 100755
index 00000000..3b2d1326
--- /dev/null
+++ b/cros_utils/device_setup_utils.py
@@ -0,0 +1,455 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Utils for setting devices
+
+This script provides utils to set device specs.
+"""
+
+from __future__ import division
+from __future__ import print_function
+
+__author__ = 'zhizhouy@google.com (Zhizhou Yang)'
+
+import re
+import time
+
+from cros_utils import command_executer
+
+
+class DutWrapper(object):
+  """Wrap DUT parameters inside.
+
+  Eventially CommandExecuter will reqiure only one
+  argument - command.
+  """
+
+  def __init__(self,
+               chromeos_root,
+               remote,
+               log_level='verbose',
+               logger=None,
+               ce=None,
+               dut_config=None):
+    self.chromeos_root = chromeos_root
+    self.remote = remote
+    self.log_level = log_level
+    self.logger = logger
+    self.ce = ce or command_executer.GetCommandExecuter(log_level=log_level)
+    self.dut_config = dut_config
+
+  def RunCommandOnDut(self, command, ignore_status=False):
+    ret, msg, err_msg = self.ce.CrosRunCommandWOutput(
+        command, machine=self.remote, chromeos_root=self.chromeos_root)
+
+    if ret:
+      err_msg = ('Command execution on DUT %s failed.\n'
+                 'Failing command: %s\n'
+                 'returned %d\n'
+                 'Error message: %s' % (self.remote, command, ret, err_msg))
+      if ignore_status:
+        self.logger.LogError(err_msg +
+                             '\n(Failure is considered non-fatal. Continue.)')
+      else:
+        self.logger.LogFatal(err_msg)
+
+    return ret, msg, err_msg
+
+  def DisableASLR(self):
+    disable_aslr = ('set -e; '
+                    'if [[ -e /proc/sys/kernel/randomize_va_space ]]; then '
+                    '  echo 0 > /proc/sys/kernel/randomize_va_space; '
+                    'fi')
+    if self.log_level == 'average':
+      self.logger.LogOutput('Disable ASLR.')
+    self.RunCommandOnDut(disable_aslr, ignore_status=False)
+
+  def SetCpuGovernor(self, governor, ignore_status=False):
+    set_gov_cmd = (
+        'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
+        # Skip writing scaling_governor if cpu is offline.
+        ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
+        '   && continue; '
+        ' cd $f; '
+        ' if [[ -e scaling_governor ]]; then '
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    if self.log_level == 'average':
+      self.logger.LogOutput('Setup CPU Governor: %s.' % governor)
+    ret, _, _ = self.RunCommandOnDut(
+        set_gov_cmd % governor, ignore_status=ignore_status)
+    return ret
+
+  def DisableTurbo(self):
+    dis_turbo_cmd = (
+        'if [[ -e /sys/devices/system/cpu/intel_pstate/no_turbo ]]; then '
+        '  if grep -q 0 /sys/devices/system/cpu/intel_pstate/no_turbo;  then '
+        '    echo -n 1 > /sys/devices/system/cpu/intel_pstate/no_turbo; '
+        '  fi; '
+        'fi; ')
+    if self.log_level == 'average':
+      self.logger.LogOutput('Disable Turbo.')
+    self.RunCommandOnDut(dis_turbo_cmd)
+
+  def SetupCpuUsage(self):
+    """Setup CPU usage.
+
+    Based on self.dut_config['cpu_usage'] configure CPU cores
+    utilization.
+    """
+
+    if (self.dut_config['cpu_usage'] == 'big_only' or
+        self.dut_config['cpu_usage'] == 'little_only'):
+      _, arch, _ = self.RunCommandOnDut('uname -m')
+
+      if arch.lower().startswith('arm') or arch.lower().startswith('aarch64'):
+        self.SetupArmCores()
+
+  def SetupArmCores(self):
+    """Setup ARM big/little cores."""
+
+    # CPU implemeters/part numbers of big/LITTLE CPU.
+    # Format: dict(CPU implementer: set(CPU part numbers))
+    LITTLE_CORES = {
+        '0x41': {
+            '0xd01',  # Cortex A32
+            '0xd03',  # Cortex A53
+            '0xd04',  # Cortex A35
+            '0xd05',  # Cortex A55
+        },
+    }
+    BIG_CORES = {
+        '0x41': {
+            '0xd07',  # Cortex A57
+            '0xd08',  # Cortex A72
+            '0xd09',  # Cortex A73
+            '0xd0a',  # Cortex A75
+            '0xd0b',  # Cortex A76
+        },
+    }
+
+    # Values of CPU Implementer and CPU part number are exposed by cpuinfo.
+    # Format:
+    # =================
+    # processor       : 0
+    # model name      : ARMv8 Processor rev 4 (v8l)
+    # BogoMIPS        : 48.00
+    # Features        : half thumb fastmult vfp edsp neon vfpv3 tls vfpv4
+    # CPU implementer : 0x41
+    # CPU architecture: 8
+    # CPU variant     : 0x0
+    # CPU part        : 0xd03
+    # CPU revision    : 4
+
+    _, cpuinfo, _ = self.RunCommandOnDut('cat /proc/cpuinfo')
+
+    # List of all CPU cores: 0, 1, ..
+    proc_matches = re.findall(r'^processor\s*: (\d+)$', cpuinfo, re.MULTILINE)
+    # List of all corresponding CPU implementers
+    impl_matches = re.findall(r'^CPU implementer\s*: (0x[\da-f]+)$', cpuinfo,
+                              re.MULTILINE)
+    # List of all corresponding CPU part numbers
+    part_matches = re.findall(r'^CPU part\s*: (0x[\da-f]+)$', cpuinfo,
+                              re.MULTILINE)
+    assert len(proc_matches) == len(impl_matches)
+    assert len(part_matches) == len(impl_matches)
+
+    all_cores = set(proc_matches)
+    dut_big_cores = {
+        core
+        for core, impl, part in zip(proc_matches, impl_matches, part_matches)
+        if impl in BIG_CORES and part in BIG_CORES[impl]
+    }
+    dut_lit_cores = {
+        core
+        for core, impl, part in zip(proc_matches, impl_matches, part_matches)
+        if impl in LITTLE_CORES and part in LITTLE_CORES[impl]
+    }
+
+    if self.dut_config['cpu_usage'] == 'big_only':
+      cores_to_enable = dut_big_cores
+      cores_to_disable = all_cores - dut_big_cores
+    elif self.dut_config['cpu_usage'] == 'little_only':
+      cores_to_enable = dut_lit_cores
+      cores_to_disable = all_cores - dut_lit_cores
+    else:
+      self.logger.LogError(
+          'cpu_usage=%s is not supported on ARM.\n'
+          'Ignore ARM CPU setup and continue.' % self.dut_config['cpu_usage'])
+      return
+
+    if cores_to_enable:
+      cmd_enable_cores = ('echo 1 | tee /sys/devices/system/cpu/cpu{%s}/online'
+                          % ','.join(sorted(cores_to_enable)))
+
+      cmd_disable_cores = ''
+      if cores_to_disable:
+        cmd_disable_cores = (
+            'echo 0 | tee /sys/devices/system/cpu/cpu{%s}/online' % ','.join(
+                sorted(cores_to_disable)))
+
+      self.RunCommandOnDut('; '.join([cmd_enable_cores, cmd_disable_cores]))
+    else:
+      # If there are no cores enabled by dut_config then configuration
+      # is invalid for current platform and should be ignored.
+      self.logger.LogError(
+          '"cpu_usage" is invalid for targeted platform.\n'
+          'dut_config[cpu_usage]=%s\n'
+          'dut big cores: %s\n'
+          'dut little cores: %s\n'
+          'Ignore ARM CPU setup and continue.' % (self.dut_config['cpu_usage'],
+                                                  dut_big_cores, dut_lit_cores))
+
+  def GetCpuOnline(self):
+    """Get online status of CPU cores.
+
+    Return dict of {int(cpu_num): <0|1>}.
+    """
+    get_cpu_online_cmd = ('paste -d" "'
+                          ' <(ls /sys/devices/system/cpu/cpu*/online)'
+                          ' <(cat /sys/devices/system/cpu/cpu*/online)')
+    _, online_output_str, _ = self.RunCommandOnDut(get_cpu_online_cmd)
+
+    # Here is the output we expect to see:
+    # -----------------
+    # /sys/devices/system/cpu/cpu0/online 0
+    # /sys/devices/system/cpu/cpu1/online 1
+
+    cpu_online = {}
+    cpu_online_match = re.compile(r'^[/\S]+/cpu(\d+)/[/\S]+\s+(\d+)$')
+    for line in online_output_str.splitlines():
+      match = cpu_online_match.match(line)
+      if match:
+        cpu = int(match.group(1))
+        status = int(match.group(2))
+        cpu_online[cpu] = status
+    # At least one CPU has to be online.
+    assert cpu_online
+
+    return cpu_online
+
+  def SetupCpuFreq(self, online_cores):
+    """Setup CPU frequency.
+
+    Based on self.dut_config['cpu_freq_pct'] setup frequency of online CPU cores
+    to a supported value which is less or equal to (freq_pct * max_freq / 100)
+    limited by min_freq.
+
+    NOTE: scaling_available_frequencies support is required.
+    Otherwise the function has no effect.
+    """
+    freq_percent = self.dut_config['cpu_freq_pct']
+    list_all_avail_freq_cmd = ('ls /sys/devices/system/cpu/cpu{%s}/cpufreq/'
+                               'scaling_available_frequencies')
+    # Ignore error to support general usage of frequency setup.
+    # Not all platforms support scaling_available_frequencies.
+    ret, all_avail_freq_str, _ = self.RunCommandOnDut(
+        list_all_avail_freq_cmd % ','.join(str(core) for core in online_cores),
+        ignore_status=True)
+    if ret or not all_avail_freq_str:
+      # No scalable frequencies available for the core.
+      return ret
+    for avail_freq_path in all_avail_freq_str.split():
+      # Get available freq from every scaling_available_frequency path.
+      # Error is considered fatal in self.RunCommandOnDut().
+      _, avail_freq_str, _ = self.RunCommandOnDut('cat ' + avail_freq_path)
+      assert avail_freq_str
+
+      all_avail_freq = sorted(
+          int(freq_str) for freq_str in avail_freq_str.split())
+      min_freq = all_avail_freq[0]
+      max_freq = all_avail_freq[-1]
+      # Calculate the frequency we are targeting.
+      target_freq = round(max_freq * freq_percent / 100)
+      # More likely it's not in the list of supported frequencies
+      # and our goal is to find the one which is less or equal.
+      # Default is min and we will try to maximize it.
+      avail_ngt_target = min_freq
+      # Find the largest not greater than the target.
+      for next_largest in reversed(all_avail_freq):
+        if next_largest <= target_freq:
+          avail_ngt_target = next_largest
+          break
+
+      max_freq_path = avail_freq_path.replace('scaling_available_frequencies',
+                                              'scaling_max_freq')
+      min_freq_path = avail_freq_path.replace('scaling_available_frequencies',
+                                              'scaling_min_freq')
+      # With default ignore_status=False we expect 0 status or Fatal error.
+      self.RunCommandOnDut('echo %s | tee %s %s' %
+                           (avail_ngt_target, max_freq_path, min_freq_path))
+
+  def WaitCooldown(self):
+    waittime = 0
+    timeout_in_sec = int(self.dut_config['cooldown_time']) * 60
+    # Temperature from sensors come in uCelsius units.
+    temp_in_ucels = int(self.dut_config['cooldown_temp']) * 1000
+    sleep_interval = 30
+
+    # Wait until any of two events occurs:
+    # 1. CPU cools down to a specified temperature.
+    # 2. Timeout cooldown_time expires.
+    # For the case when targeted temperature is not reached within specified
+    # timeout the benchmark is going to start with higher initial CPU temp.
+    # In the worst case it may affect test results but at the same time we
+    # guarantee the upper bound of waiting time.
+    # TODO(denik): Report (or highlight) "high" CPU temperature in test results.
+    # "high" should be calculated based on empirical data per platform.
+    # Based on such reports we can adjust CPU configuration or
+    # cooldown limits accordingly.
+    while waittime < timeout_in_sec:
+      _, temp_output, _ = self.RunCommandOnDut(
+          'cat /sys/class/thermal/thermal_zone*/temp', ignore_status=True)
+      if any(int(temp) > temp_in_ucels for temp in temp_output.split()):
+        time.sleep(sleep_interval)
+        waittime += sleep_interval
+      else:
+        # Exit the loop when:
+        # 1. Reported temp numbers from all thermal sensors do not exceed
+        # 'cooldown_temp' or
+        # 2. No data from the sensors.
+        break
+
+    self.logger.LogOutput('Cooldown wait time: %.1f min' % (waittime / 60))
+    return waittime
+
+  def DecreaseWaitTime(self):
+    """Change the ten seconds wait time for pagecycler to two seconds."""
+    FILE = '/usr/local/telemetry/src/tools/perf/page_sets/page_cycler_story.py'
+    ret = self.RunCommandOnDut('ls ' + FILE)
+
+    if not ret:
+      sed_command = 'sed -i "s/_TTI_WAIT_TIME = 10/_TTI_WAIT_TIME = 2/g" '
+      self.RunCommandOnDut(sed_command + FILE)
+
+  def StopUI(self):
+    self.RunCommandOnDut('stop ui')
+
+  def StartUI(self):
+    self.RunCommandOnDut('start ui')
+
+  def KerncmdUpdateNeeded(self, intel_pstate):
+    """Check whether kernel cmdline update is needed.
+
+    Args:
+      intel_pstate: kernel command line argument (active, passive, no_hwp)
+
+    Returns:
+      True if update is needed.
+    """
+
+    good = 0
+
+    # Check that dut platform supports hwp
+    cmd = "grep -q '^flags.*hwp' /proc/cpuinfo"
+    ret_code, _, _ = self.RunCommandOnDut(cmd, ignore_status=True)
+    if ret_code != good:
+      # Intel hwp is not supported, update is not needed.
+      return False
+
+    kern_cmdline_cmd = 'grep -q "intel_pstate=%s" /proc/cmdline' % intel_pstate
+    ret_code, _, _ = self.RunCommandOnDut(kern_cmdline_cmd, ignore_status=True)
+    self.logger.LogOutput('grep /proc/cmdline returned %d' % ret_code)
+    if (intel_pstate and ret_code == good or
+        not intel_pstate and ret_code != good):
+      # No need to updated cmdline if:
+      # 1. We are setting intel_pstate and we found it is already set.
+      # 2. Not using intel_pstate and it is not in cmdline.
+      return False
+
+    # Otherwise we need to update intel_pstate.
+    return True
+
+  def UpdateKerncmdIntelPstate(self, intel_pstate):
+    """Update kernel command line.
+
+    Args:
+      intel_pstate: kernel command line argument (active, passive, no_hwp)
+    """
+
+    good = 0
+
+    # First phase is to remove rootfs verification to allow cmdline change.
+    remove_verif_cmd = ' '.join([
+        '/usr/share/vboot/bin/make_dev_ssd.sh',
+        '--remove_rootfs_verification',
+        '--partition %d',
+    ])
+    # Command for partition 2.
+    verif_part2_failed, _, _ = self.RunCommandOnDut(
+        remove_verif_cmd % 2, ignore_status=True)
+    # Command for partition 4
+    # Some machines in the lab use partition 4 to boot from,
+    # so cmdline should be update for both partitions.
+    verif_part4_failed, _, _ = self.RunCommandOnDut(
+        remove_verif_cmd % 4, ignore_status=True)
+    if verif_part2_failed or verif_part4_failed:
+      self.logger.LogFatal(
+          'ERROR. Failed to update kernel cmdline on partition %d.\n'
+          'Remove verification failed with status %d' %
+          (2 if verif_part2_failed else 4, verif_part2_failed or
+           verif_part4_failed))
+
+    self.RunCommandOnDut('reboot && exit')
+    # Give enough time for dut to complete reboot
+    # TODO(denik): Replace with the function checking machine availability.
+    time.sleep(30)
+
+    # Second phase to update intel_pstate in kernel cmdline.
+    kern_cmdline = '\n'.join([
+        'tmpfile=$(mktemp)',
+        'partnumb=%d',
+        'pstate=%s',
+        # Store kernel cmdline in a temp file.
+        '/usr/share/vboot/bin/make_dev_ssd.sh --partition ${partnumb}'
+        ' --save_config ${tmpfile}',
+        # Remove intel_pstate argument if present.
+        "sed -i -r 's/ intel_pstate=[A-Za-z_]+//g' ${tmpfile}.${partnumb}",
+        # Insert intel_pstate with a new value if it is set.
+        '[[ -n ${pstate} ]] &&'
+        ' sed -i -e \"s/ *$/ intel_pstate=${pstate}/\" ${tmpfile}.${partnumb}',
+        # Save the change in kernel cmdline.
+        # After completion we have to reboot.
+        '/usr/share/vboot/bin/make_dev_ssd.sh --partition ${partnumb}'
+        ' --set_config ${tmpfile}'
+    ])
+    kern_part2_cmdline_cmd = kern_cmdline % (2, intel_pstate)
+    self.logger.LogOutput(
+        'Command to change kernel command line: %s' % kern_part2_cmdline_cmd)
+    upd_part2_failed, _, _ = self.RunCommandOnDut(
+        kern_part2_cmdline_cmd, ignore_status=True)
+    # Again here we are updating cmdline for partition 4
+    # in addition to partition 2. Without this some machines
+    # in the lab might fail.
+    kern_part4_cmdline_cmd = kern_cmdline % (4, intel_pstate)
+    self.logger.LogOutput(
+        'Command to change kernel command line: %s' % kern_part4_cmdline_cmd)
+    upd_part4_failed, _, _ = self.RunCommandOnDut(
+        kern_part4_cmdline_cmd, ignore_status=True)
+    if upd_part2_failed or upd_part4_failed:
+      self.logger.LogFatal(
+          'ERROR. Failed to update kernel cmdline on partition %d.\n'
+          'intel_pstate update failed with status %d' %
+          (2 if upd_part2_failed else 4, upd_part2_failed or upd_part4_failed))
+
+    self.RunCommandOnDut('reboot && exit')
+    # Wait 30s after reboot.
+    time.sleep(30)
+
+    # Verification phase.
+    # Check that cmdline was updated.
+    # Throw an exception if not.
+    kern_cmdline_cmd = 'grep -q "intel_pstate=%s" /proc/cmdline' % intel_pstate
+    ret_code, _, _ = self.RunCommandOnDut(kern_cmdline_cmd, ignore_status=True)
+    if (intel_pstate and ret_code != good or
+        not intel_pstate and ret_code == good):
+      # Kernel cmdline doesn't match input intel_pstate.
+      self.logger.LogFatal(
+          'ERROR. Failed to update kernel cmdline. '
+          'Final verification failed with status %d' % ret_code)
+
+    self.logger.LogOutput('Kernel cmdline updated successfully.')
diff --git a/cros_utils/device_setup_utils_unittest.py b/cros_utils/device_setup_utils_unittest.py
new file mode 100755
index 00000000..162851ed
--- /dev/null
+++ b/cros_utils/device_setup_utils_unittest.py
@@ -0,0 +1,618 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Unittest for device_setup_utils."""
+
+from __future__ import print_function
+
+import time
+
+import unittest
+import mock
+
+from device_setup_utils import DutWrapper
+
+from cros_utils import command_executer
+from cros_utils import logger
+
+BIG_LITTLE_CPUINFO = """processor       : 0
+model name      : ARMv8 Processor rev 4 (v8l)
+BogoMIPS        : 48.00
+Features        : half thumb fastmult vfp edsp neon vfpv3 tls vfpv4
+CPU implementer : 0x41
+CPU architecture: 8
+CPU variant     : 0x0
+CPU part        : 0xd03
+CPU revision    : 4
+
+processor       : 1
+model name      : ARMv8 Processor rev 4 (v8l)
+BogoMIPS        : 48.00
+Features        : half thumb fastmult vfp edsp neon vfpv3 tls vfpv4
+CPU implementer : 0x41
+CPU architecture: 8
+CPU variant     : 0x0
+CPU part        : 0xd03
+CPU revision    : 4
+
+processor       : 2
+model name      : ARMv8 Processor rev 2 (v8l)
+BogoMIPS        : 48.00
+Features        : half thumb fastmult vfp edsp neon vfpv3 tls vfpv4
+CPU implementer : 0x41
+CPU architecture: 8
+CPU variant     : 0x0
+CPU part        : 0xd08
+CPU revision    : 2
+"""
+LITTLE_ONLY_CPUINFO = """processor       : 0
+model name      : ARMv8 Processor rev 4 (v8l)
+BogoMIPS        : 48.00
+Features        : half thumb fastmult vfp edsp neon vfpv3 tls vfpv4
+CPU implementer : 0x41
+CPU architecture: 8
+CPU variant     : 0x0
+CPU part        : 0xd03
+CPU revision    : 4
+
+processor       : 1
+model name      : ARMv8 Processor rev 4 (v8l)
+BogoMIPS        : 48.00
+Features        : half thumb fastmult vfp edsp neon vfpv3 tls vfpv4
+CPU implementer : 0x41
+CPU architecture: 8
+CPU variant     : 0x0
+CPU part        : 0xd03
+CPU revision    : 4
+"""
+
+NOT_BIG_LITTLE_CPUINFO = """processor       : 0
+model name      : ARMv7 Processor rev 1 (v7l)
+Features        : swp half thumb fastmult vfp edsp thumbee neon vfpv3 tls vfpv4
+CPU implementer : 0x41
+CPU architecture: 7
+CPU variant     : 0x0
+CPU part        : 0xc0d
+CPU revision    : 1
+
+processor       : 1
+model name      : ARMv7 Processor rev 1 (v7l)
+Features        : swp half thumb fastmult vfp edsp thumbee neon vfpv3 tls vfpv4
+CPU implementer : 0x41
+CPU architecture: 7
+CPU variant     : 0x0
+CPU part        : 0xc0d
+CPU revision    : 1
+
+Hardware        : Rockchip (Device Tree)
+Revision        : 0000
+Serial          : 0000000000000000
+"""
+
+
+class DutWrapperTest(unittest.TestCase):
+  """Class of DutWrapper test."""
+  real_logger = logger.GetLogger()
+
+  mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter)
+  mock_logger = mock.Mock(spec=logger.Logger)
+
+  def __init__(self, *args, **kwargs):
+    super(DutWrapperTest, self).__init__(*args, **kwargs)
+
+  def setUp(self):
+    self.dw = DutWrapper(
+        '/tmp/chromeos',
+        'lumpy.cros2',
+        log_level='verbose',
+        logger=self.mock_logger,
+        ce=self.mock_cmd_exec,
+        dut_config={})
+
+  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
+  def test_run_command_on_dut(self, mock_cros_runcmd):
+    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd
+    mock_cros_runcmd.return_value = (0, '', '')
+    mock_cros_runcmd.assert_not_called()
+    self.dw.RunCommandOnDut('run command;')
+    mock_cros_runcmd.assert_called_once_with(
+        'run command;', chromeos_root='/tmp/chromeos', machine='lumpy.cros2')
+
+  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
+  def test_dut_wrapper_fatal_error(self, mock_cros_runcmd):
+    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd
+    # Command returns error 1.
+    mock_cros_runcmd.return_value = (1, '', 'Error!')
+    mock_cros_runcmd.assert_not_called()
+    self.dw.RunCommandOnDut('run command;')
+    mock_cros_runcmd.assert_called_once_with(
+        'run command;', chromeos_root='/tmp/chromeos', machine='lumpy.cros2')
+    # Error status causes log fatal.
+    self.assertEqual(
+        self.mock_logger.method_calls[-1],
+        mock.call.LogFatal('Command execution on DUT lumpy.cros2 failed.\n'
+                           'Failing command: run command;\nreturned 1\n'
+                           'Error message: Error!'))
+
+  @mock.patch.object(command_executer.CommandExecuter, 'CrosRunCommandWOutput')
+  def test_dut_wrapper_ignore_error(self, mock_cros_runcmd):
+    self.mock_cmd_exec.CrosRunCommandWOutput = mock_cros_runcmd
+    # Command returns error 1.
+    mock_cros_runcmd.return_value = (1, '', 'Error!')
+    self.dw.RunCommandOnDut('run command;', ignore_status=True)
+    mock_cros_runcmd.assert_called_once_with(
+        'run command;', chromeos_root='/tmp/chromeos', machine='lumpy.cros2')
+    # Error status is not fatal. LogError records the error message.
+    self.assertEqual(
+        self.mock_logger.method_calls[-1],
+        mock.call.LogError('Command execution on DUT lumpy.cros2 failed.\n'
+                           'Failing command: run command;\nreturned 1\n'
+                           'Error message: Error!\n'
+                           '(Failure is considered non-fatal. Continue.)'))
+
+  def test_disable_aslr(self):
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '', ''))
+    self.dw.DisableASLR()
+    # pyformat: disable
+    set_cpu_cmd = ('set -e; '
+                   'if [[ -e /proc/sys/kernel/randomize_va_space ]]; then '
+                   '  echo 0 > /proc/sys/kernel/randomize_va_space; '
+                   'fi')
+    self.dw.RunCommandOnDut.assert_called_once_with(
+        set_cpu_cmd, ignore_status=False)
+
+  def test_set_cpu_governor(self):
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '', ''))
+    self.dw.SetCpuGovernor('new_governor', ignore_status=False)
+    set_cpu_cmd = (
+        'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
+        # Skip writing scaling_governor if cpu is offline.
+        ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
+        '   && continue; '
+        ' cd $f; '
+        ' if [[ -e scaling_governor ]]; then '
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    self.dw.RunCommandOnDut.assert_called_once_with(
+        set_cpu_cmd % 'new_governor', ignore_status=False)
+
+  def test_set_cpu_governor_propagate_error(self):
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(1, '', 'Error.'))
+    self.dw.SetCpuGovernor('non-exist_governor')
+    set_cpu_cmd = (
+        'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
+        # Skip writing scaling_governor if cpu is not online.
+        ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
+        '   && continue; '
+        ' cd $f; '
+        ' if [[ -e scaling_governor ]]; then '
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    # By default error status is fatal.
+    self.dw.RunCommandOnDut.assert_called_once_with(
+        set_cpu_cmd % 'non-exist_governor', ignore_status=False)
+
+  def test_set_cpu_governor_ignore_status(self):
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(1, '', 'Error.'))
+    ret_code = self.dw.SetCpuGovernor('non-exist_governor', ignore_status=True)
+    set_cpu_cmd = (
+        'for f in `ls -d /sys/devices/system/cpu/cpu*/cpufreq 2>/dev/null`; do '
+        # Skip writing scaling_governor if cpu is not online.
+        ' [[ -e ${f/cpufreq/online} ]] && grep -q 0 ${f/cpufreq/online} '
+        '   && continue; '
+        ' cd $f; '
+        ' if [[ -e scaling_governor ]]; then '
+        '  echo %s > scaling_governor; fi; '
+        'done; ')
+    self.dw.RunCommandOnDut.assert_called_once_with(
+        set_cpu_cmd % 'non-exist_governor', ignore_status=True)
+    self.assertEqual(ret_code, 1)
+
+  def test_disable_turbo(self):
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '', ''))
+    self.dw.DisableTurbo()
+    set_cpu_cmd = (
+        # Disable Turbo in Intel pstate driver
+        'if [[ -e /sys/devices/system/cpu/intel_pstate/no_turbo ]]; then '
+        '  if grep -q 0 /sys/devices/system/cpu/intel_pstate/no_turbo;  then '
+        '    echo -n 1 > /sys/devices/system/cpu/intel_pstate/no_turbo; '
+        '  fi; '
+        'fi; ')
+    self.dw.RunCommandOnDut.assert_called_once_with(set_cpu_cmd)
+
+  def test_get_cpu_online_two(self):
+    """Test one digit CPU #."""
+    self.dw.RunCommandOnDut = mock.Mock(
+        return_value=(0, '/sys/devices/system/cpu/cpu0/online 0\n'
+                      '/sys/devices/system/cpu/cpu1/online 1\n', ''))
+    cpu_online = self.dw.GetCpuOnline()
+    self.assertEqual(cpu_online, {0: 0, 1: 1})
+
+  def test_get_cpu_online_twelve(self):
+    """Test two digit CPU #."""
+    self.dw.RunCommandOnDut = mock.Mock(
+        return_value=(0, '/sys/devices/system/cpu/cpu0/online 1\n'
+                      '/sys/devices/system/cpu/cpu1/online 0\n'
+                      '/sys/devices/system/cpu/cpu10/online 1\n'
+                      '/sys/devices/system/cpu/cpu11/online 1\n'
+                      '/sys/devices/system/cpu/cpu2/online 1\n'
+                      '/sys/devices/system/cpu/cpu3/online 0\n'
+                      '/sys/devices/system/cpu/cpu4/online 1\n'
+                      '/sys/devices/system/cpu/cpu5/online 0\n'
+                      '/sys/devices/system/cpu/cpu6/online 1\n'
+                      '/sys/devices/system/cpu/cpu7/online 0\n'
+                      '/sys/devices/system/cpu/cpu8/online 1\n'
+                      '/sys/devices/system/cpu/cpu9/online 0\n', ''))
+    cpu_online = self.dw.GetCpuOnline()
+    self.assertEqual(cpu_online, {
+        0: 1,
+        1: 0,
+        2: 1,
+        3: 0,
+        4: 1,
+        5: 0,
+        6: 1,
+        7: 0,
+        8: 1,
+        9: 0,
+        10: 1,
+        11: 1
+    })
+
+  def test_get_cpu_online_no_output(self):
+    """Test error case, no output."""
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '', ''))
+    with self.assertRaises(AssertionError):
+      self.dw.GetCpuOnline()
+
+  def test_get_cpu_online_command_error(self):
+    """Test error case, command error."""
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=AssertionError)
+    with self.assertRaises(AssertionError):
+      self.dw.GetCpuOnline()
+
+  @mock.patch.object(DutWrapper, 'SetupArmCores')
+  def test_setup_cpu_usage_little_on_arm(self, mock_setup_arm):
+    self.dw.SetupArmCores = mock_setup_arm
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, 'armv7l', ''))
+    self.dw.dut_config['cpu_usage'] = 'little_only'
+    self.dw.SetupCpuUsage()
+    self.dw.SetupArmCores.assert_called_once_with()
+
+  @mock.patch.object(DutWrapper, 'SetupArmCores')
+  def test_setup_cpu_usage_big_on_aarch64(self, mock_setup_arm):
+    self.dw.SetupArmCores = mock_setup_arm
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, 'aarch64', ''))
+    self.dw.dut_config['cpu_usage'] = 'big_only'
+    self.dw.SetupCpuUsage()
+    self.dw.SetupArmCores.assert_called_once_with()
+
+  @mock.patch.object(DutWrapper, 'SetupArmCores')
+  def test_setup_cpu_usage_big_on_intel(self, mock_setup_arm):
+    self.dw.SetupArmCores = mock_setup_arm
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, 'x86_64', ''))
+    self.dw.dut_config['cpu_usage'] = 'big_only'
+    self.dw.SetupCpuUsage()
+    # Check that SetupArmCores not called with invalid setup.
+    self.dw.SetupArmCores.assert_not_called()
+
+  @mock.patch.object(DutWrapper, 'SetupArmCores')
+  def test_setup_cpu_usage_all_on_intel(self, mock_setup_arm):
+    self.dw.SetupArmCores = mock_setup_arm
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, 'x86_64', ''))
+    self.dw.dut_config['cpu_usage'] = 'all'
+    self.dw.SetupCpuUsage()
+    # Check that SetupArmCores not called in general case.
+    self.dw.SetupArmCores.assert_not_called()
+
+  def test_setup_arm_cores_big_on_big_little(self):
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0, BIG_LITTLE_CPUINFO, ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_usage'] = 'big_only'
+    self.dw.SetupArmCores()
+    self.dw.RunCommandOnDut.assert_called_with(
+        'echo 1 | tee /sys/devices/system/cpu/cpu{2}/online; '
+        'echo 0 | tee /sys/devices/system/cpu/cpu{0,1}/online')
+
+  def test_setup_arm_cores_little_on_big_little(self):
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0, BIG_LITTLE_CPUINFO, ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_usage'] = 'little_only'
+    self.dw.SetupArmCores()
+    self.dw.RunCommandOnDut.assert_called_with(
+        'echo 1 | tee /sys/devices/system/cpu/cpu{0,1}/online; '
+        'echo 0 | tee /sys/devices/system/cpu/cpu{2}/online')
+
+  def test_setup_arm_cores_invalid_config(self):
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0, LITTLE_ONLY_CPUINFO, ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_usage'] = 'big_only'
+    self.dw.SetupArmCores()
+    # Check that setup command is not sent when trying
+    # to use 'big_only' on a platform with all little cores.
+    self.dw.RunCommandOnDut.assert_called_once_with('cat /proc/cpuinfo')
+
+  def test_setup_arm_cores_not_big_little(self):
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0, NOT_BIG_LITTLE_CPUINFO, ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_usage'] = 'big_only'
+    self.dw.SetupArmCores()
+    # Check that setup command is not sent when trying
+    # to use 'big_only' on a platform w/o support of big/little.
+    self.dw.RunCommandOnDut.assert_called_once_with('cat /proc/cpuinfo')
+
+  def test_setup_arm_cores_unsupported_cpu_usage(self):
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0, BIG_LITTLE_CPUINFO, ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_usage'] = 'exclusive_cores'
+    self.dw.SetupArmCores()
+    # Check that setup command is not sent when trying to use
+    # 'exclusive_cores' on ARM CPU setup.
+    self.dw.RunCommandOnDut.assert_called_once_with('cat /proc/cpuinfo')
+
+  def test_setup_cpu_freq_single_full(self):
+    online = [0]
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0,
+         '/sys/devices/system/cpu/cpu0/cpufreq/scaling_available_frequencies\n',
+         ''),
+        (0, '1 2 3 4 5 6 7 8 9 10', ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_freq_pct'] = 100
+    self.dw.SetupCpuFreq(online)
+    self.assertGreaterEqual(self.dw.RunCommandOnDut.call_count, 3)
+    self.assertEqual(
+        self.dw.RunCommandOnDut.call_args,
+        mock.call('echo 10 | tee '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_min_freq'))
+
+  def test_setup_cpu_freq_middle(self):
+    online = [0]
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0,
+         '/sys/devices/system/cpu/cpu0/cpufreq/scaling_available_frequencies\n',
+         ''),
+        (0, '1 2 3 4 5 6 7 8 9 10', ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_freq_pct'] = 60
+    self.dw.SetupCpuFreq(online)
+    self.assertGreaterEqual(self.dw.RunCommandOnDut.call_count, 2)
+    self.assertEqual(
+        self.dw.RunCommandOnDut.call_args,
+        mock.call('echo 6 | tee '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_min_freq'))
+
+  def test_setup_cpu_freq_lowest(self):
+    online = [0]
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0,
+         '/sys/devices/system/cpu/cpu0/cpufreq/scaling_available_frequencies\n',
+         ''),
+        (0, '1 2 3 4 5 6 7 8 9 10', ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_freq_pct'] = 0
+    self.dw.SetupCpuFreq(online)
+    self.assertGreaterEqual(self.dw.RunCommandOnDut.call_count, 2)
+    self.assertEqual(
+        self.dw.RunCommandOnDut.call_args,
+        mock.call('echo 1 | tee '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_min_freq'))
+
+  def test_setup_cpu_freq_multiple_middle(self):
+    online = [0, 1]
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0,
+         '/sys/devices/system/cpu/cpu0/cpufreq/scaling_available_frequencies\n'
+         '/sys/devices/system/cpu/cpu1/cpufreq/scaling_available_frequencies\n',
+         ''),
+        (0, '1 2 3 4 5 6 7 8 9 10', ''),
+        (0, '', ''),
+        (0, '1 4 6 8 10 12 14 16 18 20', ''),
+        (0, '', ''),
+    ])
+    self.dw.dut_config['cpu_freq_pct'] = 70
+    self.dw.SetupCpuFreq(online)
+    self.assertEqual(self.dw.RunCommandOnDut.call_count, 5)
+    self.assertEqual(
+        self.dw.RunCommandOnDut.call_args_list[2],
+        mock.call('echo 7 | tee '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_max_freq '
+                  '/sys/devices/system/cpu/cpu0/cpufreq/scaling_min_freq'))
+    self.assertEqual(
+        self.dw.RunCommandOnDut.call_args_list[4],
+        mock.call('echo 14 | tee '
+                  '/sys/devices/system/cpu/cpu1/cpufreq/scaling_max_freq '
+                  '/sys/devices/system/cpu/cpu1/cpufreq/scaling_min_freq'))
+
+  def test_setup_cpu_freq_no_scaling_available(self):
+    online = [0, 1]
+    self.dw.RunCommandOnDut = mock.Mock(
+        return_value=(2, '', 'No such file or directory'))
+    self.dw.dut_config['cpu_freq_pct'] = 50
+    self.dw.SetupCpuFreq(online)
+    self.dw.RunCommandOnDut.assert_called_once()
+    self.assertNotRegexpMatches(self.dw.RunCommandOnDut.call_args_list[0][0][0],
+                                '^echo.*scaling_max_freq$')
+
+  def test_setup_cpu_freq_multiple_no_access(self):
+    online = [0, 1]
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0,
+         '/sys/devices/system/cpu/cpu0/cpufreq/scaling_available_frequencies\n'
+         '/sys/devices/system/cpu/cpu1/cpufreq/scaling_available_frequencies\n',
+         ''),
+        (0, '1 4 6 8 10 12 14 16 18 20', ''),
+        AssertionError(),
+    ])
+    self.dw.dut_config['cpu_freq_pct'] = 30
+    # Error status causes log fatal.
+    with self.assertRaises(AssertionError):
+      self.dw.SetupCpuFreq(online)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_nowait(self, mock_sleep):
+    mock_sleep.return_value = 0
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '39000', ''))
+    self.dw.dut_config['cooldown_time'] = 10
+    self.dw.dut_config['cooldown_temp'] = 40
+    wait_time = self.dw.WaitCooldown()
+    # Send command to DUT only once to check temperature
+    # and make sure it does not exceed the threshold.
+    self.dw.RunCommandOnDut.assert_called_once()
+    mock_sleep.assert_not_called()
+    self.assertEqual(wait_time, 0)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_needwait_once(self, mock_sleep):
+    """Wait one iteration for cooldown.
+
+    Set large enough timeout and changing temperature
+    output. Make sure it exits when expected value
+    received.
+    Expect that WaitCooldown check temp twice.
+    """
+    mock_sleep.return_value = 0
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[(0, '41000',
+                                                      ''), (0, '39999', '')])
+    self.dw.dut_config['cooldown_time'] = 100
+    self.dw.dut_config['cooldown_temp'] = 40
+    wait_time = self.dw.WaitCooldown()
+    self.dw.RunCommandOnDut.assert_called()
+    self.assertEqual(self.dw.RunCommandOnDut.call_count, 2)
+    mock_sleep.assert_called()
+    self.assertGreater(wait_time, 0)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_needwait(self, mock_sleep):
+    """Test exit by timeout.
+
+    Send command to DUT checking the temperature and
+    check repeatedly until timeout goes off.
+    Output from temperature sensor never changes.
+    """
+    mock_sleep.return_value = 0
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '41000', ''))
+    self.dw.dut_config['cooldown_time'] = 60
+    self.dw.dut_config['cooldown_temp'] = 40
+    wait_time = self.dw.WaitCooldown()
+    self.dw.RunCommandOnDut.assert_called()
+    self.assertGreater(self.dw.RunCommandOnDut.call_count, 2)
+    mock_sleep.assert_called()
+    self.assertGreater(wait_time, 0)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_needwait_multtemp(self, mock_sleep):
+    """Wait until all temps go down.
+
+    Set large enough timeout and changing temperature
+    output. Make sure it exits when expected value
+    for all temperatures received.
+    Expect 3 checks.
+    """
+    mock_sleep.return_value = 0
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (0, '41000\n20000\n30000\n45000', ''),
+        (0, '39000\n20000\n30000\n41000', ''),
+        (0, '39000\n20000\n30000\n31000', ''),
+    ])
+    self.dw.dut_config['cooldown_time'] = 100
+    self.dw.dut_config['cooldown_temp'] = 40
+    wait_time = self.dw.WaitCooldown()
+    self.dw.RunCommandOnDut.assert_called()
+    self.assertEqual(self.dw.RunCommandOnDut.call_count, 3)
+    mock_sleep.assert_called()
+    self.assertGreater(wait_time, 0)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_thermal_error(self, mock_sleep):
+    """Handle error status.
+
+    Any error should be considered non-fatal.
+    """
+    mock_sleep.return_value = 0
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[
+        (1, '39000\n20000\n30000\n41000', 'Thermal error'),
+        (1, '39000\n20000\n30000\n31000', 'Thermal error'),
+    ])
+    self.dw.dut_config['cooldown_time'] = 10
+    self.dw.dut_config['cooldown_temp'] = 40
+    wait_time = self.dw.WaitCooldown()
+    # Check that errors are ignored.
+    self.dw.RunCommandOnDut.assert_called_with(
+        'cat /sys/class/thermal/thermal_zone*/temp', ignore_status=True)
+    self.assertEqual(self.dw.RunCommandOnDut.call_count, 2)
+    # Check that we are waiting even when an error is returned
+    # as soon as data is coming.
+    mock_sleep.assert_called()
+    self.assertGreater(wait_time, 0)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_thermal_no_output(self, mock_sleep):
+    """Handle no output.
+
+    Check handling of empty stdout.
+    """
+    mock_sleep.return_value = 0
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[(1, '', 'Thermal error')])
+    self.dw.dut_config['cooldown_time'] = 10
+    self.dw.dut_config['cooldown_temp'] = 40
+    wait_time = self.dw.WaitCooldown()
+    # Check that errors are ignored.
+    self.dw.RunCommandOnDut.assert_called_once_with(
+        'cat /sys/class/thermal/thermal_zone*/temp', ignore_status=True)
+    # No wait.
+    mock_sleep.assert_not_called()
+    self.assertEqual(wait_time, 0)
+
+  @mock.patch.object(time, 'sleep')
+  def test_wait_cooldown_thermal_ws_output(self, mock_sleep):
+    """Handle whitespace output.
+
+    Check handling of whitespace only.
+    """
+    mock_sleep.return_value = 0
+    self.dw.RunCommandOnDut = mock.Mock(side_effect=[(1, '\n',
+                                                      'Thermal error')])
+    self.dw.dut_config['cooldown_time'] = 10
+    self.dw.dut_config['cooldown_temp'] = 40
+    wait_time = self.dw.WaitCooldown()
+    # Check that errors are ignored.
+    self.dw.RunCommandOnDut.assert_called_once_with(
+        'cat /sys/class/thermal/thermal_zone*/temp', ignore_status=True)
+    # No wait.
+    mock_sleep.assert_not_called()
+    self.assertEqual(wait_time, 0)
+
+  def test_stop_ui(self):
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '', ''))
+    self.dw.StopUI()
+    self.dw.RunCommandOnDut.assert_called_once_with('stop ui')
+
+  def test_start_ui(self):
+    self.dw.RunCommandOnDut = mock.Mock(return_value=(0, '', ''))
+    self.dw.StartUI()
+    self.dw.RunCommandOnDut.assert_called_once_with('start ui')
+
+
+if __name__ == '__main__':
+  unittest.main()