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 If the given images all have no serial associated and have the same 248 image for the same board, testbed will assign all devices with the 249 desired board to the image. This allows tests to randomly pick devices 250 to run. 251 As an example, a testbed with 4 devices, 2 for board_1 and 2 for 252 board_2. If the given images value is: 253 [('board_1_build', None), ('board_2_build', None)] 254 The testbed will return following device allocation: 255 {'serial_1_board_1': 'board_1_build', 256 'serial_2_board_1': 'board_1_build', 257 'serial_1_board_2': 'board_2_build', 258 'serial_2_board_2': 'board_2_build', 259 } 260 That way, all board_1 duts will be installed with board_1_build, and 261 all board_2 duts will be installed with board_2_build. Test can pick 262 any dut from board_1 duts and same applies to board_2 duts. 263 264 @param images: A list of tuples of (build, serial). serial could be None 265 if it's not specified. Following are some examples: 266 [('branch1/shamu-userdebug/100', None), 267 ('branch1/shamu-userdebug/100', None)] 268 [('branch1/hammerhead-userdebug/100', 'XZ123'), 269 ('branch1/hammerhead-userdebug/200', None)] 270 where XZ123 is serial of one of the hammerheads connected to the 271 testbed. 272 273 @return: A dictionary of (serial, build). Note that build here should 274 not have a serial specified in it. 275 @raise InstallError: If not enough duts are available to install the 276 given images. Or there are more duts with the same board than 277 the images list specified. 278 """ 279 # The map between serial and build to install in that dut. 280 serial_build_pairs = {} 281 builds_without_serial = [build for build, serial in images 282 if not serial] 283 for build, serial in images: 284 if serial: 285 serial_build_pairs[serial] = build 286 # Return the mapping if all builds have serial specified. 287 if not builds_without_serial: 288 return serial_build_pairs 289 290 # serials grouped by the board of duts. 291 duts_by_name = {} 292 for serial, host in self.get_adb_devices().iteritems(): 293 # Excluding duts already assigned to a build. 294 if serial in serial_build_pairs: 295 continue 296 aliases = host.get_device_aliases() 297 for alias in aliases: 298 duts_by_name.setdefault(alias, []).append(serial) 299 300 # Builds grouped by the board name. 301 builds_by_name = {} 302 for build in builds_without_serial: 303 match = re.match(adb_host.BUILD_REGEX, build) 304 if not match: 305 raise error.InstallError('Build %s is invalid. Failed to parse ' 306 'the board name.' % build) 307 name = match.group('BUILD_TARGET') 308 builds_by_name.setdefault(name, []).append(build) 309 310 # Pair build with dut with matching board. 311 for name, builds in builds_by_name.iteritems(): 312 duts = duts_by_name.get(name, []) 313 if len(duts) < len(builds): 314 raise error.InstallError( 315 'Expected number of DUTs for name %s is %d, got %d' % 316 (name, len(builds), len(duts) if duts else 0)) 317 elif len(duts) == len(builds): 318 serial_build_pairs.update(dict(zip(duts, builds))) 319 else: 320 # In this cases, available dut number is greater than the number 321 # of builds. 322 if len(set(builds)) > 1: 323 raise error.InstallError( 324 'Number of available DUTs are greater than builds ' 325 'needed, testbed cannot allocate DUTs for testing ' 326 'deterministically.') 327 # Set all DUTs to the same build. 328 for serial in duts: 329 serial_build_pairs[serial] = builds[0] 330 331 return serial_build_pairs 332 333 334 def save_info(self, results_dir): 335 """Saves info about the testbed to a directory. 336 337 @param results_dir: The directory to save to. 338 """ 339 for device in self.get_adb_devices().values(): 340 device.save_info(results_dir, include_build_info=True) 341 342 343 def _stage_shared_build(self, serial_build_map): 344 """Try to stage build on teststation to be shared by all provision jobs. 345 346 This logic only applies to the case that multiple devices are 347 provisioned to the same build. If the provision job does not fit this 348 requirement, this method will not stage any build. 349 350 @param serial_build_map: A map between dut's serial and the build to be 351 installed. 352 353 @return: A tuple of (build_url, build_local_path, teststation), where 354 build_url: url to the build on devserver 355 build_local_path: Path to a local directory in teststation that 356 contains the build. 357 teststation: A teststation object that is used to stage the 358 build. 359 If there are more than one build need to be staged or only one 360 device is used for the test, return (None, None, None) 361 """ 362 build_local_path = None 363 build_url = None 364 teststation = None 365 same_builds = set([build for build in serial_build_map.values()]) 366 if len(same_builds) == 1 and len(serial_build_map.values()) > 1: 367 same_build = same_builds.pop() 368 logging.debug('All devices will be installed with build %s, stage ' 369 'the shared build to be used for all provision jobs.', 370 same_build) 371 stage_host = self.get_adb_devices()[serial_build_map.keys()[0]] 372 teststation = stage_host.teststation 373 build_url, _ = stage_host.stage_build_for_install(same_build) 374 if stage_host.get_os_type() == adb_host.OS_TYPE_ANDROID: 375 build_local_path = stage_host.stage_android_image_files( 376 build_url) 377 else: 378 build_local_path = stage_host.stage_brillo_image_files( 379 build_url) 380 elif len(same_builds) > 1: 381 logging.debug('More than one build need to be staged, leave the ' 382 'staging build tasks to individual provision job.') 383 else: 384 logging.debug('Only one device needs to be provisioned, leave the ' 385 'staging build task to individual provision job.') 386 387 return build_url, build_local_path, teststation 388 389 390 def machine_install(self, image=None): 391 """Install the DUT. 392 393 @param image: Image we want to install on this testbed, e.g., 394 `branch1/shamu-eng/1001,branch2/shamu-eng/1002` 395 396 @returns A tuple of (the name of the image installed, None), where None 397 is a placeholder for update_url. Testbed does not have a single 398 update_url, thus it's set to None. 399 @returns A tuple of (image_name, host_attributes). 400 image_name is the name of images installed, e.g., 401 `branch1/shamu-eng/1001,branch2/shamu-eng/1002` 402 host_attributes is a dictionary of (attribute, value), which 403 can be saved to afe_host_attributes table in database. This 404 method returns a dictionary with entries of job_repo_urls for 405 each provisioned devices: 406 `job_repo_url_[adb_serial]`: devserver_url, where devserver_url 407 is a url to the build staged on devserver. 408 For example: 409 {'job_repo_url_XZ001': 'http://10.1.1.3/branch1/shamu-eng/1001', 410 'job_repo_url_XZ002': 'http://10.1.1.3/branch2/shamu-eng/1002'} 411 """ 412 image = image or self._parser.options.image 413 if not image: 414 raise error.InstallError('No image string is provided to test bed.') 415 images = self._parse_image(image) 416 host_attributes = {} 417 418 # Change logging formatter to include thread name. This is to help logs 419 # from each provision runs have the dut's serial, which is set as the 420 # thread name. 421 logging_config.add_threadname_in_log() 422 423 serial_build_map = self.locate_devices(images) 424 425 build_url, build_local_path, teststation = self._stage_shared_build( 426 serial_build_map) 427 428 thread_pool = None 429 try: 430 arguments = [] 431 for serial, build in serial_build_map.iteritems(): 432 logging.info('Installing build %s on DUT with serial %s.', 433 build, serial) 434 host = self.get_adb_devices()[serial] 435 if build_url: 436 device_build_url = build_url 437 else: 438 device_build_url, _ = host.stage_build_for_install(build) 439 arguments.append({'host': host, 440 'build_url': device_build_url, 441 'build_local_path': build_local_path}) 442 attribute_name = '%s_%s' % (constants.JOB_REPO_URL, 443 host.adb_serial) 444 host_attributes[attribute_name] = device_build_url 445 446 thread_pool = pool.ThreadPool(_POOL_SIZE) 447 thread_pool.map(self._install_device, arguments) 448 thread_pool.close() 449 except Exception as err: 450 logging.error(err.message) 451 finally: 452 if thread_pool: 453 thread_pool.join() 454 455 if build_local_path: 456 logging.debug('Clean up build artifacts %s:%s', 457 teststation.hostname, build_local_path) 458 teststation.run('rm -rf %s' % build_local_path) 459 460 return image, host_attributes 461 462 463 def get_attributes_to_clear_before_provision(self): 464 """Get a list of attribute to clear before machine_install starts. 465 """ 466 return [host.job_repo_url_attribute for host in 467 self.adb_devices.values()] 468