write_firmware: Support select of a particular servo

Sometimes it is useful to have multiple servo boards connected. Allow a
particular board to be selected for dut-control with a --servo option.

BUG=chromium-os:34729
TEST=manual

Try different options manually to make sure they work:

My snow servo is on 9999, link on 8888. Check that everything works as
expected:

Snow:

(cros) $ cros_bundle_firmware -b daisy --servo 9999 -w usb

Snow: (which is the default on my machine)

(cros) $ cros_bundle_firmware -b daisy --servo any -w usb

Gives an error:
    raise IOError('No servo access available, please use --servo')
(cros) $ cros_bundle_firmware -b daisy --servo none -w usb

Link:
(cros) $ cros_bundle_firmware -b link --servo 8888 -w em100

Change-Id: I9ab9b043c5d01eb7ab2efaace666cdcc22013e89
Signed-off-by: Simon Glass <sjg@chromium.org>
Reviewed-on: https://gerrit.chromium.org/gerrit/33968
Reviewed-by: Randall Spangler <rspangler@chromium.org>
diff --git a/host/lib/write_firmware.py b/host/lib/write_firmware.py
index 66c152c..158bd77 100644
--- a/host/lib/write_firmware.py
+++ b/host/lib/write_firmware.py
@@ -35,6 +35,13 @@
   and instructions to flash it to SPI flash. The payload is itself normally a
   full Chrome OS image consisting of U-Boot, some keys and verification
   information, images and a map of the flash memory.
+
+  Private attributes:
+    _servo_port: Port number to use to talk to servo with dut-control.
+      Special values are:
+        None: servo is not available.
+        0: any servo will do.
+
   """
   def __init__(self, tools, fdt, output, bundle):
     """Set up a new WriteFirmware object.
@@ -55,6 +62,26 @@
     self.update = True
     self.verify = False
 
+    # Use default servo port
+    self._servo_port = 0
+
+  def SelectServo(self, servo):
+    """Select the servo to use for writing firmware.
+
+    Args:
+      servo: String containing description of servo to use:
+        'none'  : Don't use servo, generate an error on any attempt.
+        'any'   : Use any available servo.
+        '<port>': Use servo with that port number.
+    """
+    if servo == 'none':
+      self._servo_port = None
+    elif servo == 'any':
+      self._servo_port = 0
+    else:
+      self._servo_port = int(servo)
+    self._out.Notice('Servo port %s' % str(self._servo_port))
+
   def _GetFlashScript(self, payload_size, update, verify, boot_type, checksum,
                       bus='0'):
     """Get the U-Boot boot command needed to flash U-Boot.
@@ -325,6 +352,23 @@
 
     return False
 
+  def _DutControl(self, args):
+    """Run dut-control with supplied arguments.
+
+    The correct servo will be used based on self._servo_port.
+
+    Args:
+      args: List of arguments to dut-control.
+
+    Raises:
+      IOError if no servo access is permitted.
+    """
+    if self._servo_port is None:
+      raise IOError('No servo access available, please use --servo')
+    if self._servo_port:
+      args.extend(['-p', '%s' % self._servo_port])
+    self._tools.Run('dut-control', args)
+
   def _ExtractPayloadParts(self, payload):
     """Extract the BL1, BL2 and U-Boot parts from a payload.
 
@@ -395,7 +439,7 @@
     # back to life.
     # BUG=chromium-os:28229
     args = ['cold_reset:on', 'sleep:.2', 'cold_reset:off'] + args
-    self._tools.Run('dut-control', args)
+    self._DutControl(args)
 
     # If we have a kernel to write, create a new image with that added.
     if kernel:
@@ -436,12 +480,12 @@
         if upto == 1:
           # Once SPL starts up we can release the power buttom
           args = ['fw_up:off', 'pwr_button:release']
-          self._tools.Run('dut-control', args)
+          self._DutControl(args)
 
     finally:
       # Make sure that the power button is released, whatever happens
       args = ['fw_up:off', 'pwr_button:release']
-      self._tools.Run('dut-control', args)
+      self._DutControl(args)
 
     self._out.Notice('Image downloaded - please see serial output '
         'for progress.')
@@ -605,7 +649,7 @@
     """
     args = ['spi2_vref:off', 'spi2_buf_en:off', 'spi2_buf_on_flex_en:off']
     args.append('spi_hold:on')
-    self._tools.Run('dut-control', args)
+    self._DutControl(args)
 
     # TODO(sjg@chromium.org): This is for link. We could make this
     # configurable from the fdt.
@@ -616,12 +660,12 @@
     self._out.Progress('Resetting board')
     args = ['cold_reset:on', 'sleep:.2', 'cold_reset:off', 'sleep:.5']
     args.extend(['pwr_button:press', 'sleep:.2', 'pwr_button:release'])
-    self._tools.Run('dut-control', args)
+    self._DutControl(args)
 
 
 def DoWriteFirmware(output, tools, fdt, flasher, file_list, image_fname,
                     bundle, update=True, verify=False, dest=None,
-                    flash_dest=None, kernel=None, props={}):
+                    flash_dest=None, kernel=None, props={}, servo='any'):
   """A simple function to write firmware to a device.
 
   This creates a WriteFirmware object and uses it to write the firmware image
@@ -641,8 +685,11 @@
     flash_dest: Destination device for flasher to program payload into.
     kernel: Kernel file to write after U-Boot
     props: A dictionary containing properties from the PackFirmware object
+    servo: Describes the servo unit to use: none=none; any=any; otherwise
+           port number of servo to use.
   """
   write = WriteFirmware(tools, fdt, output, bundle)
+  write.SelectServo(servo)
   write.update = update
   write.verify = verify
   if dest == 'usb':