tigertool: fix reboot

Fix command truncation issue and change reboot detection mechanism.
Fix lint.

BUG=b:216199797
TEST=./tigertest.py -c 10
BRANCH=none

Signed-off-by: Nick Sanders <nsanders@chromium.org>
Change-Id: I2a3f962c2bcc1d87a2c3e8ac7da89e99997f3b40
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/3419581
Reviewed-by: Ruben Rodriguez Buchillon <coconutruben@chromium.org>

5 files changed, 228 insertions, 30 deletions

diff --git a/extra/tigertool/ecusb/pty_driver.py b/extra/tigertool/ecusb/pty_driver.py
index 137cbc149b..09ef8c42e4 100644
--- a/extra/tigertool/ecusb/pty_driver.py
+++ b/extra/tigertool/ecusb/pty_driver.py
@@ -171,6 +171,9 @@ class ptyDriver(object):
     except pexpect.TIMEOUT:
       raise ptyError('Timeout waiting for response.')
     finally:
+      if not regex_list:
+        # Must be longer than delaybeforesend
+        time.sleep(0.1)
       self._close()
     return result_list
 
diff --git a/extra/tigertool/ecusb/tiny_servo_common.py b/extra/tigertool/ecusb/tiny_servo_common.py
index 152c238bdf..e27736a9dc 100644
--- a/extra/tigertool/ecusb/tiny_servo_common.py
+++ b/extra/tigertool/ecusb/tiny_servo_common.py
@@ -17,6 +17,7 @@ import re
 import subprocess
 import sys
 import time
+import usb
 
 import six
 
@@ -86,6 +87,69 @@ def check_usb_sn(vidpid):
 
   return None
 
+def get_usb_dev(vidpid, serialname=None):
+  """Return the USB pyusb devie struct
+
+  Return the dev struct of the first USB device with VID:PID vidpid,
+  or None if no device is found. If more than one device check serial
+  if supplied.
+
+  Args:
+    vidpid: string representation of the usb vid:pid, eg. '18d1:2001'
+    serialname: serialname if specified.
+
+  Returns: pyusb device if found, None otherwise.
+  """
+  vidpidst = vidpid.split(':')
+  vid = int(vidpidst[0], 16)
+  pid = int(vidpidst[1], 16)
+
+
+  dev_g = usb.core.find(idVendor=vid, idProduct=pid, find_all=True)
+  dev_list = list(dev_g)
+
+  if not dev_list:
+    return None
+
+  # Check if we have multiple devices and we've specified the serial.
+  dev = None
+  if serialname:
+    for d in dev_list:
+      dev_serial = usb.util.get_string(d, d.iSerialNumber)
+      if dev_serial == serialname:
+        dev = d
+        break
+    if dev is None:
+      return None
+  else:
+    try:
+      dev = dev_list[0]
+    except StopIteration:
+      return None
+
+  return dev
+
+def check_usb_dev(vidpid, serialname=None):
+  """Return the USB dev number
+
+  Return the dev number of the first USB device with VID:PID vidpid,
+  or None if no device is found. If more than one device check serial
+  if supplied.
+
+  Args:
+    vidpid: string representation of the usb vid:pid, eg. '18d1:2001'
+    serialname: serialname if specified.
+
+  Returns: usb device number if found, None otherwise.
+  """
+  dev = get_usb_dev(vidpid, serialname=serialname)
+
+  if dev:
+    return dev.address
+
+  return None
+
+
 def wait_for_usb_remove(vidpid, serialname=None, timeout=None):
   """Wait for USB device with vidpid to be removed.
 
@@ -109,7 +173,7 @@ def wait_for_usb(vidpid, serialname=None, timeout=None, desiredpresence=True):
   if timeout:
     finish = datetime.datetime.now() + datetime.timedelta(seconds=timeout)
   while check_usb(vidpid, serialname) != desiredpresence:
-    time.sleep(.1)
+    time.sleep(.01)
     if timeout:
       if datetime.datetime.now() > finish:
         raise TinyServoError('Timeout', 'Timeout waiting for USB %s' % vidpid)
diff --git a/extra/tigertool/make_pkg.sh b/extra/tigertool/make_pkg.sh
index d2860f64c1..5a63862242 100755
--- a/extra/tigertool/make_pkg.sh
+++ b/extra/tigertool/make_pkg.sh
@@ -17,6 +17,8 @@ cp ../usb_serial/console.py "${DEST}"
 cp ../../../../../chroot/usr/bin/dfu-util "${DEST}"
 cp flash_dfu.sh "${DEST}"
 cp tigertool.py "${DEST}"
+cp tigertest.py "${DEST}"
+cp README.md "${DEST}"
 
 cp -r ecusb "${DEST}"
 cp -r ../../../../../chroot/usr/lib64/python2.7/site-packages/usb "${DEST}"
diff --git a/extra/tigertool/tigertest.py b/extra/tigertool/tigertest.py
new file mode 100755
index 0000000000..0cd31c8cce
--- /dev/null
+++ b/extra/tigertool/tigertest.py
@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+# Copyright 2022 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.
+#
+# Ignore indention messages, since legacy scripts use 2 spaces instead of 4.
+# pylint: disable=bad-indentation,docstring-section-indent
+# pylint: disable=docstring-trailing-quotes
+
+"""Smoke test of tigertool binary."""
+
+import argparse
+import subprocess
+import sys
+
+
+# Script to control tigertail USB-C Mux board.
+#
+# optional arguments:
+#  -h, --help            show this help message and exit
+#  -s SERIALNO, --serialno SERIALNO
+#                        serial number of board to use
+#  -b BUS, --bus BUS     Which rail to log: [vbus|cc1|cc2]
+#  --setserialno SETSERIALNO
+#                        serial number to set on the board.
+#  --check_serial        check serial number set on the board.
+#  -m MUX, --mux MUX     mux selection
+#  -p, --power           check VBUS
+#  -l POWERLOG, --powerlog POWERLOG
+#                        log VBUS
+#  -r SYSJUMP, --sysjump SYSJUMP
+#                        region selection
+#  --reboot              reboot tigertail
+#  --check_version       check tigertail version
+
+
+def testCmd(cmd, expected_results):
+  """Run command on console, check for success.
+
+  Args:
+    cmd: shell command to run.
+    expected_results: a list object of strings expected in the result.
+
+  Raises:
+    Exception on fail.
+  """
+  print('run: ' + cmd)
+  try:
+    p = subprocess.run(cmd, shell=True, check=False, capture_output=True)
+    output = p.stdout.decode('utf-8')
+    error = p.stderr.decode('utf-8')
+    assert p.returncode == 0
+    for result in expected_results:
+      output.index(result)
+  except Exception as e:
+    print('FAIL')
+    print('cmd: ' + cmd)
+    print('error: ' + str(e))
+    print('stdout:\n' + output)
+    print('stderr:\n' + error)
+    print('expected: ' + str(expected_results))
+    print('RC: ' + str(p.returncode))
+    raise e
+
+def test_sequence():
+  testCmd('./tigertool.py --reboot', ['PASS'])
+  testCmd('./tigertool.py --setserialno test', ['PASS'])
+  testCmd('./tigertool.py --check_serial', ['test', 'PASS'])
+  testCmd('./tigertool.py -s test --check_serial', ['test', 'PASS'])
+  testCmd('./tigertool.py -m A', ['Mux set to A', 'PASS'])
+  testCmd('./tigertool.py -m B', ['Mux set to B', 'PASS'])
+  testCmd('./tigertool.py -m off', ['Mux set to off', 'PASS'])
+  testCmd('./tigertool.py -p', ['PASS'])
+  testCmd('./tigertool.py -r rw', ['PASS'])
+  testCmd('./tigertool.py -r ro', ['PASS'])
+  testCmd('./tigertool.py --check_version', ['RW', 'RO', 'PASS'])
+
+  print('PASS')
+
+def main(argv):
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument('-c', '--count', type=int, default=1,
+                      help='loops to run')
+
+  opts = parser.parse_args(argv)
+
+  for i in range(1, opts.count + 1):
+    print('Iteration: %d' % i)
+    test_sequence()
+
+if __name__ == '__main__':
+    main(sys.argv[1:])
diff --git a/extra/tigertool/tigertool.py b/extra/tigertool/tigertool.py
index 79aa30c3a4..6baae8abdf 100755
--- a/extra/tigertool/tigertool.py
+++ b/extra/tigertool/tigertool.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 # Copyright 2017 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.
@@ -36,7 +36,7 @@ def do_mux(mux, pty):
     c.log('Mux setting %s invalid, try one of %s' % (mux, validmux))
     return False
 
-  cmd = '\r\nmux %s\r\n' % mux
+  cmd = 'mux %s' % mux
   regex = 'TYPE\-C mux is ([^\s\r\n]*)\r'
 
   results = pty._issue_cmd_get_results(cmd, [regex])[0]
@@ -62,11 +62,10 @@ def do_version(pty):
   # RW:      tigertail_v1.1.6749-74d1a312e
   # Build:   tigertail_v1.1.6749-74d1a312e
   #          2017-07-25 20:08:34 nsanders@meatball.mtv.corp.google.com
-
   """
-  cmd = '\r\nversion\r\n'
-  regex = 'RO:\s+(\S+)\s+RW:\s+(\S+)\s+Build:\s+(\S+)\s+' \
-          '(\d\d\d\d-\d\d-\d\d \d\d:\d\d:\d\d) (\S+)'
+  cmd = 'version'
+  regex = r'RO:\s+(\S+)\s+RW:\s+(\S+)\s+Build:\s+(\S+)\s+' \
+          r'(\d\d\d\d-\d\d-\d\d \d\d:\d\d:\d\d) (\S+)'
 
   results = pty._issue_cmd_get_results(cmd, [regex])[0]
   c.log('Version is %s' % results[3])
@@ -87,8 +86,8 @@ def do_check_serial(pty):
   # > serialno
   # Serial number: number
   """
-  cmd = '\r\nserialno\r\n'
-  regex = 'Serial number: ([^\n\r]+)'
+  cmd = 'serialno'
+  regex = r'Serial number: ([^\n\r]+)'
 
   results = pty._issue_cmd_get_results(cmd, [regex])[0]
   c.log('Serial is %s' % results[1])
@@ -126,45 +125,71 @@ def do_power(count, bus, pty):
 
   c.log('time,\tmV,\tmW,\tmA')
 
-  cmd = '\r\nina %s\r\n' % ina
-  regex = 'Bus voltage  : \S+ \S+ (\d+) mV\s+' \
-          'Power        : \S+ \S+ (\d+) mW\s+' \
-          'Current      : \S+ \S+ (\d+) mA'
+  cmd = 'ina %s' % ina
+  regex = r'Bus voltage  : \S+ \S+ (\d+) mV\s+' \
+          r'Power        : \S+ \S+ (\d+) mW\s+' \
+          r'Current      : \S+ \S+ (\d+) mA'
 
   for i in range(0, count):
     results = pty._issue_cmd_get_results(cmd, [regex])[0]
-    c.log('%.2f,\t%s,\t%s\t%s' % (time.time() - start,
-          results[1], results[2], results[3]))
+    c.log('%.2f,\t%s,\t%s\t%s' % (
+        time.time() - start,
+        results[1], results[2], results[3]))
 
   return True
 
-def do_reboot(pty):
+def do_reboot(pty, serialname):
   """Reboot via ec console pty
 
   Args:
     pty: a pty object connected to tigertail
+    serialname: serial name, can be None.
 
   Command is: reboot.
   """
-  cmd = '\r\nreboot\r\n'
-  regex = 'Rebooting'
+  cmd = 'reboot'
+
+  # Check usb dev number on current instance.
+  devno = c.check_usb_dev(STM_VIDPID, serialname=serialname)
+  if not devno:
+    c.log('Device not found')
+    return False
 
   try:
-    results = pty._issue_cmd_get_results(cmd, [regex])[0]
-    time.sleep(1)
-    c.log(results)
+    pty._issue_cmd(cmd)
   except Exception as e:
-    c.log(e)
+    c.log('Failed to send command: ' + str(e))
+    return False
+
+  try:
+    c.wait_for_usb_remove(STM_VIDPID, timeout=3., serialname=serialname)
+  except Exception as e:
+    # Polling for reboot isn't reliable but if it hasn't happened in 3 seconds
+    # it's not going to. This step just goes faster if it's detected.
+    pass
+
+  try:
+    c.wait_for_usb(STM_VIDPID, timeout=3., serialname=serialname)
+  except Exception as e:
+    c.log('Failed to return from reboot: ' + str(e))
+    return False
+
+  # Check that the device had a new device number, i.e. it's
+  # disconnected and reconnected.
+  newdevno = c.check_usb_dev(STM_VIDPID, serialname=serialname)
+  if newdevno == devno:
+    c.log("Device didn't reboot")
     return False
 
   return True
 
-def do_sysjump(region, pty):
+def do_sysjump(region, pty, serialname):
   """Set region via ec console 'pty'.
 
   Args:
     region: ec code region to execute, 'ro' or 'rw'
     pty: a pty object connected to tigertail
+    serialname: serial name, can be None.
 
   Commands are:
   # > sysjump rw
@@ -175,12 +200,24 @@ def do_sysjump(region, pty):
         region, validregion))
     return False
 
-  cmd = '\r\nsysjump %s\r\n' % region
+  cmd = 'sysjump %s' % region
   try:
     pty._issue_cmd(cmd)
-    time.sleep(1)
   except Exception as e:
-    c.log(e)
+    c.log('Exception: ' + str(e))
+    return False
+
+  try:
+    c.wait_for_usb_remove(STM_VIDPID, timeout=3., serialname=serialname)
+  except Exception as e:
+    # Polling for reboot isn't reliable but if it hasn't happened in 3 seconds
+    # it's not going to. This step just goes faster if it's detected.
+    pass
+
+  try:
+    c.wait_for_usb(STM_VIDPID, timeout=3., serialname=serialname)
+  except Exception as e:
+    c.log('Failed to return from restart: ' + str(e))
     return False
 
   c.log('Region requested %s' % region)
@@ -192,7 +229,7 @@ def get_parser():
   parser.add_argument('-s', '--serialno', type=str, default=None,
                       help='serial number of board to use')
   parser.add_argument('-b', '--bus', type=str, default='vbus',
-                     help='Which rail to log: [vbus|cc1|cc2]')
+                      help='Which rail to log: [vbus|cc1|cc2]')
   group = parser.add_mutually_exclusive_group()
   group.add_argument('--setserialno', type=str, default=None,
                      help='serial number to set on the board.')
@@ -238,10 +275,10 @@ def main(argv):
     result &= do_mux(opts.mux, pty)
 
   elif opts.sysjump:
-    result &= do_sysjump(opts.sysjump, pty)
+    result &= do_sysjump(opts.sysjump, pty, serialname=opts.serialno)
 
   elif opts.reboot:
-    result &= do_reboot(pty)
+    result &= do_reboot(pty, serialname=opts.serialno)
 
   elif opts.check_version:
     result &= do_version(pty)
@@ -259,7 +296,7 @@ def main(argv):
     c.log('PASS')
   else:
     c.log('FAIL')
-    exit(-1)
+    sys.exit(-1)
 
 
 if __name__ == '__main__':