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