1# Copyright 2015 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5"""This class defines the TestBed class.""" 6 7import logging 8import re 9import sys 10import threading 11import traceback 12from multiprocessing import pool 13 14import common 15 16from autotest_lib.client.common_lib import error 17from autotest_lib.client.common_lib import logging_config 18from autotest_lib.server.cros.dynamic_suite import constants 19from autotest_lib.server import autoserv_parser 20from autotest_lib.server import utils 21from autotest_lib.server.cros import provision 22from autotest_lib.server.hosts import adb_host 23from autotest_lib.server.hosts import base_label 24from autotest_lib.server.hosts import host_info 25from autotest_lib.server.hosts import testbed_label 26from autotest_lib.server.hosts import teststation_host 27 28 29# Thread pool size to provision multiple devices in parallel. 30_POOL_SIZE = 4 31 32# Pattern for the image name when used to provision a dut connected to testbed. 33# It should follow the naming convention of 34# branch/target/build_id[:serial][#count], 35# where serial and count are optional. Count is the number of devices to 36# provision to. 37_IMAGE_NAME_PATTERN = '(.*/.*/[^:#]*)(?::(.*))?(?:#(\d+))?' 38 39class TestBed(object): 40 """This class represents a collection of connected teststations and duts.""" 41 42 _parser = autoserv_parser.autoserv_parser 43 VERSION_PREFIX = provision.TESTBED_BUILD_VERSION_PREFIX 44 support_devserver_provision = False 45 46 def __init__(self, hostname='localhost', afe_host=None, adb_serials=None, 47 host_info_store=None, **dargs): 48 """Initialize a TestBed. 49 50 This will create the Test Station Host and connected hosts (ADBHost for 51 now) and allow the user to retrieve them. 52 53 @param hostname: Hostname of the test station connected to the duts. 54 @param adb_serials: List of adb device serials. 55 @param host_info_store: A CachingHostInfoStore object. 56 @param afe_host: The host object attained from the AFE (get_hosts). 57 """ 58 logging.info('Initializing TestBed centered on host: %s', hostname) 59 self.hostname = hostname 60 self._afe_host = afe_host or utils.EmptyAFEHost() 61 self.host_info_store = (host_info_store or 62 host_info.InMemoryHostInfoStore()) 63 self.labels = base_label.LabelRetriever(testbed_label.TESTBED_LABELS) 64 self.teststation = teststation_host.create_teststationhost( 65 hostname=hostname, afe_host=self._afe_host, **dargs) 66 self.is_client_install_supported = False 67 serials_from_attributes = self._afe_host.attributes.get('serials') 68 if serials_from_attributes: 69 serials_from_attributes = serials_from_attributes.split(',') 70 71 self.adb_device_serials = (adb_serials or 72 serials_from_attributes or 73 self.query_adb_device_serials()) 74 self.adb_devices = {} 75 for adb_serial in self.adb_device_serials: 76 self.adb_devices[adb_serial] = adb_host.ADBHost( 77 hostname=hostname, teststation=self.teststation, 78 adb_serial=adb_serial, afe_host=self._afe_host, 79 host_info_store=self.host_info_store, **dargs) 80 81 82 def query_adb_device_serials(self): 83 """Get a list of devices currently attached to the test station. 84 85 @returns a list of adb devices. 86 """ 87 return adb_host.ADBHost.parse_device_serials( 88 self.teststation.run('adb devices').stdout) 89 90 91 def get_all_hosts(self): 92 """Return a list of all the hosts in this testbed. 93 94 @return: List of the hosts which includes the test station and the adb 95 devices. 96 """ 97 device_list = [self.teststation] 98 device_list.extend(self.adb_devices.values()) 99 return device_list 100 101 102 def get_test_station(self): 103 """Return the test station host object. 104 105 @return: The test station host object. 106 """ 107 return self.teststation 108 109 110 def get_adb_devices(self): 111 """Return the adb host objects. 112 113 @return: A dict of adb device serials to their host objects. 114 """ 115 return self.adb_devices 116 117 118 def get_labels(self): 119 """Return a list of the labels gathered from the devices connected. 120 121 @return: A list of strings that denote the labels from all the devices 122 connected. 123 """ 124 return self.labels.get_labels(self) 125 126 127 def update_labels(self): 128 """Update the labels on the testbed.""" 129 return self.labels.update_labels(self) 130 131 132 def get_platform(self): 133 """Return the platform of the devices. 134 135 @return: A string representing the testbed platform. 136 """ 137 return 'testbed' 138 139 140 def repair(self): 141 """Run through repair on all the devices.""" 142 # board name is needed for adb_host to repair as the adb_host objects 143 # created for testbed doesn't have host label and attributes retrieved 144 # from AFE. 145 info = self.host_info_store.get() 146 board = info.board 147 # Remove the tailing -# in board name as it can be passed in from 148 # testbed board labels 149 match = re.match(r'^(.*)-\d+$', board) 150 if match: 151 board = match.group(1) 152 failures = [] 153 for adb_device in self.get_adb_devices().values(): 154 try: 155 adb_device.repair(board=board, os=info.os) 156 except: 157 exc_type, exc_value, exc_traceback = sys.exc_info() 158 failures.append((adb_device.adb_serial, exc_type, exc_value, 159 exc_traceback)) 160 if failures: 161 serials = [] 162 for serial, exc_type, exc_value, exc_traceback in failures: 163 serials.append(serial) 164 details = ''.join(traceback.format_exception( 165 exc_type, exc_value, exc_traceback)) 166 logging.error('Failed to repair device with serial %s, ' 167 'error:\n%s', serial, details) 168 raise error.AutoservRepairTotalFailure( 169 'Fail to repair %d devices: %s' % 170 (len(serials), ','.join(serials))) 171 172 173 def verify(self): 174 """Run through verify on all the devices.""" 175 for device in self.get_all_hosts(): 176 device.verify() 177 178 179 def cleanup(self): 180 """Run through cleanup on all the devices.""" 181 for adb_device in self.get_adb_devices().values(): 182 adb_device.cleanup() 183 184 185 def _parse_image(self, image_string): 186 """Parse the image string to a dictionary. 187 188 Sample value of image_string: 189 Provision dut with serial ZX1G2 to build `branch1/shamu-userdebug/111`, 190 and provision another shamu with build `branch2/shamu-userdebug/222` 191 branch1/shamu-userdebug/111:ZX1G2,branch2/shamu-userdebug/222 192 193 Provision 10 shamu with build `branch1/shamu-userdebug/LATEST` 194 branch1/shamu-userdebug/LATEST#10 195 196 @param image_string: A comma separated string of images. The image name 197 is in the format of branch/target/build_id[:serial]. Serial is 198 optional once testbed machine_install supports allocating DUT 199 based on board. 200 201 @returns: A list of tuples of (build, serial). serial could be None if 202 it's not specified. 203 """ 204 images = [] 205 for image in image_string.split(','): 206 match = re.match(_IMAGE_NAME_PATTERN, image) 207 # The image string cannot specify both serial and count. 208 if not match or (match.group(2) and match.group(3)): 209 raise error.InstallError( 210 'Image name of "%s" has invalid format. It should ' 211 'follow naming convention of ' 212 'branch/target/build_id[:serial][#count]', image) 213 if match.group(3): 214 images.extend([(match.group(1), None)]*int(match.group(3))) 215 else: 216 images.append((match.group(1), match.group(2))) 217 return images 218 219 220 @staticmethod 221 def _install_device(inputs): 222 """Install build to a device with the given inputs. 223 224 @param inputs: A dictionary of the arguments needed to install a device. 225 Keys include: 226 host: An ADBHost object of the device. 227 build_url: Devserver URL to the build to install. 228 """ 229 host = inputs['host'] 230 build_url = inputs['build_url'] 231 build_local_path = inputs['build_local_path'] 232 233 # Set the thread name with the serial so logging for installing 234 # different devices can have different thread name. 235 threading.current_thread().name = host.adb_serial 236 logging.info('Starting installing device %s:%s from build url %s', 237 host.hostname, host.adb_serial, build_url) 238 host.machine_install(build_url=build_url, 239 build_local_path=build_local_path) 240 logging.info('Finished installing device %s:%s from build url %s', 241 host.hostname, host.adb_serial, build_url) 242 243 244 def locate_devices(self, images): 245 """Locate device for each image in the given images list. 246 247 @param images: A list of tuples of (build, serial). serial could be None 248 if it's not specified. Following are some examples: 249 [('branch1/shamu-userdebug/100', None), 250 ('branch1/shamu-userdebug/100', None)] 251 [('branch1/hammerhead-userdebug/100', 'XZ123'), 252 ('branch1/hammerhead-userdebug/200', None)] 253 where XZ123 is serial of one of the hammerheads connected to the 254 testbed. 255 256 @return: A dictionary of (serial, build). Note that build here should 257 not have a serial specified in it. 258 @raise InstallError: If not enough duts are available to install the 259 given images. Or there are more duts with the same board than 260 the images list specified. 261 """ 262 # The map between serial and build to install in that dut. 263 serial_build_pairs = {} 264 builds_without_serial = [build for build, serial in images 265 if not serial] 266 for build, serial in images: 267 if serial: 268 serial_build_pairs[serial] = build 269 # Return the mapping if all builds have serial specified. 270 if not builds_without_serial: 271 return serial_build_pairs 272 273 # serials grouped by the board of duts. 274 duts_by_name = {} 275 for serial, host in self.get_adb_devices().iteritems(): 276 # Excluding duts already assigned to a build. 277 if serial in serial_build_pairs: 278 continue 279 aliases = host.get_device_aliases() 280 for alias in aliases: 281 duts_by_name.setdefault(alias, []).append(serial) 282 283 # Builds grouped by the board name. 284 builds_by_name = {} 285 for build in builds_without_serial: 286 match = re.match(adb_host.BUILD_REGEX, build) 287 if not match: 288 raise error.InstallError('Build %s is invalid. Failed to parse ' 289 'the board name.' % build) 290 name = match.group('BUILD_TARGET') 291 builds_by_name.setdefault(name, []).append(build) 292 293 # Pair build with dut with matching board. 294 for name, builds in builds_by_name.iteritems(): 295 duts = duts_by_name.get(name, []) 296 if len(duts) != len(builds): 297 raise error.InstallError( 298 'Expected number of DUTs for name %s is %d, got %d' % 299 (name, len(builds), len(duts) if duts else 0)) 300 serial_build_pairs.update(dict(zip(duts, builds))) 301 return serial_build_pairs 302 303 304 def save_info(self, results_dir): 305 """Saves info about the testbed to a directory. 306 307 @param results_dir: The directory to save to. 308 """ 309 for device in self.get_adb_devices().values(): 310 device.save_info(results_dir, include_build_info=True) 311 312 313 def _stage_shared_build(self, serial_build_map): 314 """Try to stage build on teststation to be shared by all provision jobs. 315 316 This logic only applies to the case that multiple devices are 317 provisioned to the same build. If the provision job does not fit this 318 requirement, this method will not stage any build. 319 320 @param serial_build_map: A map between dut's serial and the build to be 321 installed. 322 323 @return: A tuple of (build_url, build_local_path, teststation), where 324 build_url: url to the build on devserver 325 build_local_path: Path to a local directory in teststation that 326 contains the build. 327 teststation: A teststation object that is used to stage the 328 build. 329 If there are more than one build need to be staged or only one 330 device is used for the test, return (None, None, None) 331 """ 332 build_local_path = None 333 build_url = None 334 teststation = None 335 same_builds = set([build for build in serial_build_map.values()]) 336 if len(same_builds) == 1 and len(serial_build_map.values()) > 1: 337 same_build = same_builds.pop() 338 logging.debug('All devices will be installed with build %s, stage ' 339 'the shared build to be used for all provision jobs.', 340 same_build) 341 stage_host = self.get_adb_devices()[serial_build_map.keys()[0]] 342 teststation = stage_host.teststation 343 build_url, _ = stage_host.stage_build_for_install(same_build) 344 if stage_host.get_os_type() == adb_host.OS_TYPE_ANDROID: 345 build_local_path = stage_host.stage_android_image_files( 346 build_url) 347 else: 348 build_local_path = stage_host.stage_brillo_image_files( 349 build_url) 350 elif len(same_builds) > 1: 351 logging.debug('More than one build need to be staged, leave the ' 352 'staging build tasks to individual provision job.') 353 else: 354 logging.debug('Only one device needs to be provisioned, leave the ' 355 'staging build task to individual provision job.') 356 357 return build_url, build_local_path, teststation 358 359 360 def machine_install(self, image=None): 361 """Install the DUT. 362 363 @param image: Image we want to install on this testbed, e.g., 364 `branch1/shamu-eng/1001,branch2/shamu-eng/1002` 365 366 @returns A tuple of (the name of the image installed, None), where None 367 is a placeholder for update_url. Testbed does not have a single 368 update_url, thus it's set to None. 369 @returns A tuple of (image_name, host_attributes). 370 image_name is the name of images installed, e.g., 371 `branch1/shamu-eng/1001,branch2/shamu-eng/1002` 372 host_attributes is a dictionary of (attribute, value), which 373 can be saved to afe_host_attributes table in database. This 374 method returns a dictionary with entries of job_repo_urls for 375 each provisioned devices: 376 `job_repo_url_[adb_serial]`: devserver_url, where devserver_url 377 is a url to the build staged on devserver. 378 For example: 379 {'job_repo_url_XZ001': 'http://10.1.1.3/branch1/shamu-eng/1001', 380 'job_repo_url_XZ002': 'http://10.1.1.3/branch2/shamu-eng/1002'} 381 """ 382 image = image or self._parser.options.image 383 if not image: 384 raise error.InstallError('No image string is provided to test bed.') 385 images = self._parse_image(image) 386 host_attributes = {} 387 388 # Change logging formatter to include thread name. This is to help logs 389 # from each provision runs have the dut's serial, which is set as the 390 # thread name. 391 logging_config.add_threadname_in_log() 392 393 serial_build_map = self.locate_devices(images) 394 395 build_url, build_local_path, teststation = self._stage_shared_build( 396 serial_build_map) 397 398 try: 399 arguments = [] 400 for serial, build in serial_build_map.iteritems(): 401 logging.info('Installing build %s on DUT with serial %s.', 402 build, serial) 403 host = self.get_adb_devices()[serial] 404 if build_url: 405 device_build_url = build_url 406 else: 407 device_build_url, _ = host.stage_build_for_install(build) 408 arguments.append({'host': host, 409 'build_url': device_build_url, 410 'build_local_path': build_local_path}) 411 attribute_name = '%s_%s' % (constants.JOB_REPO_URL, 412 host.adb_serial) 413 host_attributes[attribute_name] = device_build_url 414 415 thread_pool = pool.ThreadPool(_POOL_SIZE) 416 thread_pool.map(self._install_device, arguments) 417 thread_pool.close() 418 finally: 419 if build_local_path: 420 logging.debug('Clean up build artifacts %s:%s', 421 teststation.hostname, build_local_path) 422 teststation.run('rm -rf %s' % build_local_path) 423 424 return image, host_attributes 425 426 427 def get_attributes_to_clear_before_provision(self): 428 """Get a list of attribute to clear before machine_install starts. 429 """ 430 return [host.job_repo_url_attribute for host in 431 self.adb_devices.values()] 432