# Copyright 2015 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.
"""This class defines the TestBed class."""
import logging
import re
from multiprocessing import pool
import common
from autotest_lib.client.common_lib import error
from autotest_lib.server.cros.dynamic_suite import constants
from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
from autotest_lib.server import autoserv_parser
from autotest_lib.server.hosts import adb_host
from autotest_lib.server.hosts import teststation_host
# Thread pool size to provision multiple devices in parallel.
_POOL_SIZE = 4
# Pattern for the image name when used to provision a dut connected to testbed.
# It should follow the naming convention of branch/target/build_id[:serial],
# where serial is optional.
_IMAGE_NAME_PATTERN = '(.*/.*/[^:]*)(?::(.*))?'
class TestBed(object):
"""This class represents a collection of connected teststations and duts."""
_parser = autoserv_parser.autoserv_parser
VERSION_PREFIX = 'testbed-version'
def __init__(self, hostname='localhost', host_attributes={},
adb_serials=None, **dargs):
"""Initialize a TestBed.
This will create the Test Station Host and connected hosts (ADBHost for
now) and allow the user to retrieve them.
@param hostname: Hostname of the test station connected to the duts.
@param serials: List of adb device serials.
"""
logging.info('Initializing TestBed centered on host: %s', hostname)
self.hostname = hostname
self.teststation = teststation_host.create_teststationhost(
hostname=hostname)
self.is_client_install_supported = False
serials_from_attributes = host_attributes.get('serials')
if serials_from_attributes:
serials_from_attributes = serials_from_attributes.split(',')
self.adb_device_serials = (adb_serials or
serials_from_attributes or
self.query_adb_device_serials())
self.adb_devices = {}
for adb_serial in self.adb_device_serials:
self.adb_devices[adb_serial] = adb_host.ADBHost(
hostname=hostname, teststation=self.teststation,
adb_serial=adb_serial)
def query_adb_device_serials(self):
"""Get a list of devices currently attached to the test station.
@returns a list of adb devices.
"""
serials = []
# Let's see if we can get the serials via host attributes.
afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
serials_attr = afe.get_host_attribute('serials', hostname=self.hostname)
for serial_attr in serials_attr:
serials.extend(serial_attr.value.split(','))
# Looks like we got nothing from afe, let's probe the test station.
if not serials:
# TODO(kevcheng): Refactor teststation to be a class and make the
# ADBHost adb_devices a static method I can use here. For now this
# is pretty much a c/p of the _adb_devices() method from ADBHost.
serials = adb_host.ADBHost.parse_device_serials(
self.teststation.run('adb devices').stdout)
return serials
def get_all_hosts(self):
"""Return a list of all the hosts in this testbed.
@return: List of the hosts which includes the test station and the adb
devices.
"""
device_list = [self.teststation]
device_list.extend(self.adb_devices.values())
return device_list
def get_test_station(self):
"""Return the test station host object.
@return: The test station host object.
"""
return self.teststation
def get_adb_devices(self):
"""Return the adb host objects.
@return: A dict of adb device serials to their host objects.
"""
return self.adb_devices
def get_labels(self):
"""Return a list of the labels gathered from the devices connected.
@return: A list of strings that denote the labels from all the devices
connected.
"""
labels = []
for adb_device in self.get_adb_devices().values():
labels.extend(adb_device.get_labels())
# Currently the board label will need to be modified for each adb
# device. We'll get something like 'board:android-shamu' and
# we'll need to update it to 'board:android-shamu-1'. Let's store all
# the labels in a dict and keep track of how many times we encounter
# it, that way we know what number to append.
board_label_dict = {}
updated_labels = []
for label in labels:
# Update the board labels
if label.startswith(constants.BOARD_PREFIX):
# Now let's grab the board num and append it to the board_label.
board_num = board_label_dict.setdefault(label, 0) + 1
board_label_dict[label] = board_num
updated_labels.append('%s-%d' % (label, board_num))
else:
# We don't need to mess with this.
updated_labels.append(label)
return updated_labels
def get_platform(self):
"""Return the platform of the devices.
@return: A string representing the testbed platform.
"""
return 'testbed'
def repair(self):
"""Run through repair on all the devices."""
for adb_device in self.get_adb_devices().values():
adb_device.repair()
def verify(self):
"""Run through verify on all the devices."""
for device in self.get_all_hosts():
device.verify()
def cleanup(self):
"""Run through cleanup on all the devices."""
for adb_device in self.get_adb_devices().values():
adb_device.cleanup()
def _parse_image(self, image_string):
"""Parse the image string to a dictionary.
Sample value of image_string:
branch1/shamu-userdebug/LATEST:ZX1G2,branch2/shamu-userdebug/LATEST
@param image_string: A comma separated string of images. The image name
is in the format of branch/target/build_id[:serial]. Serial is
optional once testbed machine_install supports allocating DUT
based on board.
@returns: A list of tuples of (build, serial). serial could be None if
it's not specified.
"""
images = []
for image in image_string.split(','):
match = re.match(_IMAGE_NAME_PATTERN, image)
if not match:
raise error.InstallError(
'Image name of "%s" has invalid format. It should '
'follow naming convention of '
'branch/target/build_id[:serial]', image)
images.append((match.group(1), match.group(2)))
return images
@staticmethod
def _install_device(inputs):
"""Install build to a device with the given inputs.
@param inputs: A dictionary of the arguments needed to install a device.
Keys include:
host: An ADBHost object of the device.
build_url: Devserver URL to the build to install.
"""
host = inputs['host']
build_url = inputs['build_url']
logging.info('Starting installing device %s:%s from build url %s',
host.hostname, host.adb_serial, build_url)
host.machine_install(build_url=build_url)
logging.info('Finished installing device %s:%s from build url %s',
host.hostname, host.adb_serial, build_url)
def locate_devices(self, images):
"""Locate device for each image in the given images list.
@param images: A list of tuples of (build, serial). serial could be None
if it's not specified. Following are some examples:
[('branch1/shamu-userdebug/100', None),
('branch1/shamu-userdebug/100', None)]
[('branch1/hammerhead-userdebug/100', 'XZ123'),
('branch1/hammerhead-userdebug/200', None)]
where XZ123 is serial of one of the hammerheads connected to the
testbed.
@return: A dictionary of (serial, build). Note that build here should
not have a serial specified in it.
@raise InstallError: If not enough duts are available to install the
given images. Or there are more duts with the same board than
the images list specified.
"""
# The map between serial and build to install in that dut.
serial_build_pairs = {}
builds_without_serial = [build for build, serial in images
if not serial]
for build, serial in images:
if serial:
serial_build_pairs[serial] = build
# Return the mapping if all builds have serial specified.
if not builds_without_serial:
return serial_build_pairs
# serials grouped by the board of duts.
duts_by_board = {}
for serial, host in self.get_adb_devices().iteritems():
# Excluding duts already assigned to a build.
if serial in serial_build_pairs:
continue
board = host.get_board_name()
duts_by_board.setdefault(board, []).append(serial)
# Builds grouped by the board name.
builds_by_board = {}
for build in builds_without_serial:
match = re.match(adb_host.BUILD_REGEX, build)
if not match:
raise error.InstallError('Build %s is invalid. Failed to parse '
'the board name.' % build)
board = match.group('BOARD')
builds_by_board.setdefault(board, []).append(build)
# Pair build with dut with matching board.
for board, builds in builds_by_board.iteritems():
duts = duts_by_board.get(board, None)
if not duts or len(duts) != len(builds):
raise error.InstallError(
'Expected number of DUTs for board %s is %d, got %d' %
(board, len(builds), len(duts) if duts else 0))
serial_build_pairs.update(dict(zip(duts, builds)))
return serial_build_pairs
def machine_install(self):
"""Install the DUT.
@returns The name of the image installed.
"""
if not self._parser.options.image:
raise error.InstallError('No image string is provided to test bed.')
images = self._parse_image(self._parser.options.image)
arguments = []
for serial, build in self.locate_devices(images).iteritems():
logging.info('Installing build %s on DUT with serial %s.', build,
serial)
host = self.get_adb_devices()[serial]
build_url, _ = host.stage_build_for_install(build)
arguments.append({'host': host,
'build_url': build_url})
thread_pool = pool.ThreadPool(_POOL_SIZE)
thread_pool.map(self._install_device, arguments)
thread_pool.close()
return self._parser.options.image