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 9from multiprocessing import pool 10 11import common 12 13from autotest_lib.client.common_lib import error 14from autotest_lib.server.cros.dynamic_suite import constants 15from autotest_lib.server.cros.dynamic_suite import frontend_wrappers 16from autotest_lib.server import autoserv_parser 17from autotest_lib.server.hosts import adb_host 18from autotest_lib.server.hosts import teststation_host 19 20 21# Thread pool size to provision multiple devices in parallel. 22_POOL_SIZE = 4 23 24# Pattern for the image name when used to provision a dut connected to testbed. 25# It should follow the naming convention of branch/target/build_id[:serial], 26# where serial is optional. 27_IMAGE_NAME_PATTERN = '(.*/.*/[^:]*)(?::(.*))?' 28 29class TestBed(object): 30 """This class represents a collection of connected teststations and duts.""" 31 32 _parser = autoserv_parser.autoserv_parser 33 VERSION_PREFIX = 'testbed-version' 34 35 def __init__(self, hostname='localhost', host_attributes={}, 36 adb_serials=None, **dargs): 37 """Initialize a TestBed. 38 39 This will create the Test Station Host and connected hosts (ADBHost for 40 now) and allow the user to retrieve them. 41 42 @param hostname: Hostname of the test station connected to the duts. 43 @param serials: List of adb device serials. 44 """ 45 logging.info('Initializing TestBed centered on host: %s', hostname) 46 self.hostname = hostname 47 self.teststation = teststation_host.create_teststationhost( 48 hostname=hostname) 49 self.is_client_install_supported = False 50 serials_from_attributes = host_attributes.get('serials') 51 if serials_from_attributes: 52 serials_from_attributes = serials_from_attributes.split(',') 53 54 self.adb_device_serials = (adb_serials or 55 serials_from_attributes or 56 self.query_adb_device_serials()) 57 self.adb_devices = {} 58 for adb_serial in self.adb_device_serials: 59 self.adb_devices[adb_serial] = adb_host.ADBHost( 60 hostname=hostname, teststation=self.teststation, 61 adb_serial=adb_serial) 62 63 64 def query_adb_device_serials(self): 65 """Get a list of devices currently attached to the test station. 66 67 @returns a list of adb devices. 68 """ 69 serials = [] 70 # Let's see if we can get the serials via host attributes. 71 afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10) 72 serials_attr = afe.get_host_attribute('serials', hostname=self.hostname) 73 for serial_attr in serials_attr: 74 serials.extend(serial_attr.value.split(',')) 75 76 # Looks like we got nothing from afe, let's probe the test station. 77 if not serials: 78 # TODO(kevcheng): Refactor teststation to be a class and make the 79 # ADBHost adb_devices a static method I can use here. For now this 80 # is pretty much a c/p of the _adb_devices() method from ADBHost. 81 serials = adb_host.ADBHost.parse_device_serials( 82 self.teststation.run('adb devices').stdout) 83 84 return serials 85 86 87 def get_all_hosts(self): 88 """Return a list of all the hosts in this testbed. 89 90 @return: List of the hosts which includes the test station and the adb 91 devices. 92 """ 93 device_list = [self.teststation] 94 device_list.extend(self.adb_devices.values()) 95 return device_list 96 97 98 def get_test_station(self): 99 """Return the test station host object. 100 101 @return: The test station host object. 102 """ 103 return self.teststation 104 105 106 def get_adb_devices(self): 107 """Return the adb host objects. 108 109 @return: A dict of adb device serials to their host objects. 110 """ 111 return self.adb_devices 112 113 114 def get_labels(self): 115 """Return a list of the labels gathered from the devices connected. 116 117 @return: A list of strings that denote the labels from all the devices 118 connected. 119 """ 120 labels = [] 121 for adb_device in self.get_adb_devices().values(): 122 labels.extend(adb_device.get_labels()) 123 # Currently the board label will need to be modified for each adb 124 # device. We'll get something like 'board:android-shamu' and 125 # we'll need to update it to 'board:android-shamu-1'. Let's store all 126 # the labels in a dict and keep track of how many times we encounter 127 # it, that way we know what number to append. 128 board_label_dict = {} 129 updated_labels = [] 130 for label in labels: 131 # Update the board labels 132 if label.startswith(constants.BOARD_PREFIX): 133 # Now let's grab the board num and append it to the board_label. 134 board_num = board_label_dict.setdefault(label, 0) + 1 135 board_label_dict[label] = board_num 136 updated_labels.append('%s-%d' % (label, board_num)) 137 else: 138 # We don't need to mess with this. 139 updated_labels.append(label) 140 return updated_labels 141 142 143 def get_platform(self): 144 """Return the platform of the devices. 145 146 @return: A string representing the testbed platform. 147 """ 148 return 'testbed' 149 150 151 def repair(self): 152 """Run through repair on all the devices.""" 153 for adb_device in self.get_adb_devices().values(): 154 adb_device.repair() 155 156 157 def verify(self): 158 """Run through verify on all the devices.""" 159 for device in self.get_all_hosts(): 160 device.verify() 161 162 163 def cleanup(self): 164 """Run through cleanup on all the devices.""" 165 for adb_device in self.get_adb_devices().values(): 166 adb_device.cleanup() 167 168 169 def _parse_image(self, image_string): 170 """Parse the image string to a dictionary. 171 172 Sample value of image_string: 173 branch1/shamu-userdebug/LATEST:ZX1G2,branch2/shamu-userdebug/LATEST 174 175 @param image_string: A comma separated string of images. The image name 176 is in the format of branch/target/build_id[:serial]. Serial is 177 optional once testbed machine_install supports allocating DUT 178 based on board. 179 180 @returns: A list of tuples of (build, serial). serial could be None if 181 it's not specified. 182 """ 183 images = [] 184 for image in image_string.split(','): 185 match = re.match(_IMAGE_NAME_PATTERN, image) 186 if not match: 187 raise error.InstallError( 188 'Image name of "%s" has invalid format. It should ' 189 'follow naming convention of ' 190 'branch/target/build_id[:serial]', image) 191 images.append((match.group(1), match.group(2))) 192 return images 193 194 195 @staticmethod 196 def _install_device(inputs): 197 """Install build to a device with the given inputs. 198 199 @param inputs: A dictionary of the arguments needed to install a device. 200 Keys include: 201 host: An ADBHost object of the device. 202 build_url: Devserver URL to the build to install. 203 """ 204 host = inputs['host'] 205 build_url = inputs['build_url'] 206 207 logging.info('Starting installing device %s:%s from build url %s', 208 host.hostname, host.adb_serial, build_url) 209 host.machine_install(build_url=build_url) 210 logging.info('Finished installing device %s:%s from build url %s', 211 host.hostname, host.adb_serial, build_url) 212 213 214 def locate_devices(self, images): 215 """Locate device for each image in the given images list. 216 217 @param images: A list of tuples of (build, serial). serial could be None 218 if it's not specified. Following are some examples: 219 [('branch1/shamu-userdebug/100', None), 220 ('branch1/shamu-userdebug/100', None)] 221 [('branch1/hammerhead-userdebug/100', 'XZ123'), 222 ('branch1/hammerhead-userdebug/200', None)] 223 where XZ123 is serial of one of the hammerheads connected to the 224 testbed. 225 226 @return: A dictionary of (serial, build). Note that build here should 227 not have a serial specified in it. 228 @raise InstallError: If not enough duts are available to install the 229 given images. Or there are more duts with the same board than 230 the images list specified. 231 """ 232 # The map between serial and build to install in that dut. 233 serial_build_pairs = {} 234 builds_without_serial = [build for build, serial in images 235 if not serial] 236 for build, serial in images: 237 if serial: 238 serial_build_pairs[serial] = build 239 # Return the mapping if all builds have serial specified. 240 if not builds_without_serial: 241 return serial_build_pairs 242 243 # serials grouped by the board of duts. 244 duts_by_board = {} 245 for serial, host in self.get_adb_devices().iteritems(): 246 # Excluding duts already assigned to a build. 247 if serial in serial_build_pairs: 248 continue 249 board = host.get_board_name() 250 duts_by_board.setdefault(board, []).append(serial) 251 252 # Builds grouped by the board name. 253 builds_by_board = {} 254 for build in builds_without_serial: 255 match = re.match(adb_host.BUILD_REGEX, build) 256 if not match: 257 raise error.InstallError('Build %s is invalid. Failed to parse ' 258 'the board name.' % build) 259 board = match.group('BOARD') 260 builds_by_board.setdefault(board, []).append(build) 261 262 # Pair build with dut with matching board. 263 for board, builds in builds_by_board.iteritems(): 264 duts = duts_by_board.get(board, None) 265 if not duts or len(duts) != len(builds): 266 raise error.InstallError( 267 'Expected number of DUTs for board %s is %d, got %d' % 268 (board, len(builds), len(duts) if duts else 0)) 269 serial_build_pairs.update(dict(zip(duts, builds))) 270 return serial_build_pairs 271 272 273 def machine_install(self): 274 """Install the DUT. 275 276 @returns The name of the image installed. 277 """ 278 if not self._parser.options.image: 279 raise error.InstallError('No image string is provided to test bed.') 280 images = self._parse_image(self._parser.options.image) 281 282 arguments = [] 283 for serial, build in self.locate_devices(images).iteritems(): 284 logging.info('Installing build %s on DUT with serial %s.', build, 285 serial) 286 host = self.get_adb_devices()[serial] 287 build_url, _ = host.stage_build_for_install(build) 288 arguments.append({'host': host, 289 'build_url': build_url}) 290 291 thread_pool = pool.ThreadPool(_POOL_SIZE) 292 thread_pool.map(self._install_device, arguments) 293 thread_pool.close() 294 return self._parser.options.image 295