294 lines
11 KiB
Python
294 lines
11 KiB
Python
# 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
|