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.