servo_updater: enable paths to work on test images as well

This change adds some logic to find config and firmware files on test
images as well, where they end up being located at /usr/local/share
instead.

BUG=b:120921028
BRANCH=None
TEST=manual testing // on fizz-labstation test image
sudo servo_updater --board servo_v4

observe servo v4 getting updated

Note: Test images right now do not have usb_updater2 yet (WIP) hence the
test is done using fizz-labstation, a board that does come with
usb_updater2

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

1 files changed, 34 insertions, 4 deletions

diff --git a/extra/usb_updater/servo_updater.py b/extra/usb_updater/servo_updater.py
index 68a920ec1b..2a01ea7069 100755
--- a/extra/usb_updater/servo_updater.py
+++ b/extra/usb_updater/servo_updater.py
@@ -17,12 +17,33 @@ import json
 import fw_update
 import ecusb.tiny_servo_common as c
 
-FIRMWARE_PATH = "/usr/share/servo_updater/firmware/"
-CONFIGS_PATH = "/usr/share/servo_updater/configs/"
 
 class ServoUpdaterException(Exception):
   """Raised on exceptions generated by servo_updater."""
 
+
+DEFAULT_BASE_PATH = '/usr/'
+TEST_IMAGE_BASE_PATH = '/usr/local/'
+
+COMMON_PATH = 'share/servo_updater'
+
+FIRMWARE_DIR = "firmware/"
+CONFIGS_DIR = "configs/"
+
+if os.path.exists(os.path.join(DEFAULT_BASE_PATH, COMMON_PATH)):
+  BASE_PATH = DEFAULT_BASE_PATH
+elif os.path.exists(os.path.join(TEST_IMAGE_BASE_PATH, COMMON_PATH)):
+  BASE_PATH = TEST_IMAGE_BASE_PATH
+else:
+  raise ServoUpdaterException('servo_updater/ dir not found in known spots.')
+
+FIRMWARE_PATH = os.path.join(BASE_PATH, COMMON_PATH, FIRMWARE_DIR)
+CONFIGS_PATH = os.path.join(BASE_PATH, COMMON_PATH, CONFIGS_DIR)
+
+for p in [FIRMWARE_PATH, CONFIGS_PATH]:
+  if not os.path.exists(p):
+    raise ServoUpdaterException('Could not find required path %r' % p)
+
 def flash(brdfile, serialno, binfile):
   """Call fw_update to upload to updater USB endpoint."""
   p = fw_update.Supdate()
@@ -42,19 +63,28 @@ def flash(brdfile, serialno, binfile):
 
 def flash2(vidpid, serialno, binfile):
   """Call fw update via usb_updater2 commandline."""
-  cmd = "usb_updater2 -d %s" % vidpid
+  tool = 'usb_updater2'
+  cmd = "%s -d %s" % (tool, vidpid)
   if serialno:
     cmd += " -S %s" % serialno
   cmd += " -n"
   cmd += " %s" % binfile
 
   print(cmd)
+  help_cmd = '%s --help' % tool
+  with open('/dev/null') as devnull:
+    sanity_check = subprocess.call(help_cmd.split(), stdout=devnull,
+                                   stderr=devnull)
+  if sanity_check:
+    raise ServoUpdaterException('%s exit with res = %d. Make sure the tool '
+                                'is available on the device.' % (help_cmd,
+                                                                 sanity_check))
   res = subprocess.call(cmd.split())
 
   if res in (0, 1, 2):
     return res
   else:
-    raise ServoUpdaterException("usb_updater2 exit with res = %d" % res)
+    raise ServoUpdaterException("%s exit with res = %d" % (cmd, res))
 
 def connect(vidpid, iface, serialno, debuglog=False):
   """Connect to console.