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
5import logging
6import multiprocessing
7import sys
8
9from autotest_lib.client.common_lib import error
10from autotest_lib.client.common_lib import global_config
11from autotest_lib.client.common_lib.cros import dev_server
12from autotest_lib.server.cros import autoupdater
13from autotest_lib.server.cros.dynamic_suite import constants
14
15#Update status
16UPDATE_SUCCESS = 0
17UPDATE_FAILURE = 1
18
19def update_dut_worker(updater_obj, dut, image, force):
20    """The method called by multiprocessing worker pool for updating DUT.
21    This function is the function which is repeatedly scheduled for each
22    DUT through the multiprocessing worker. This has to be defined outside
23    the class because it needs to be pickleable.
24
25    @param updater_obj: An CliqueDUTUpdater object.
26    @param dut: DUTObject representing the DUT.
27    @param image: The build type and version to install on the host.
28    @param force: If False, will only updated the host if it is not
29                  already running the build. If True, force the
30                  update regardless, and force a full-reimage.
31
32    """
33    updater_obj.update_dut(dut_host=dut.host, image=image, force=force)
34
35
36class CliqueDUTUpdater(object):
37    """CliqueDUTUpdater is responsible for updating all the DUT's in the
38    DUT pool to the same release.
39    """
40
41    def __init__(self):
42        """Initializes the DUT updater for updating the DUT's in the pool."""
43
44
45    @staticmethod
46    def _get_board_name_from_host(dut_host):
47        """Get the board name of the remote host.
48
49        @param host: Host object representing the DUT.
50
51        @return: A string representing the board of the remote host.
52        """
53        try:
54            board = dut_host.get_board().replace(constants.BOARD_PREFIX, '')
55        except error.AutoservRunError:
56            raise error.TestFail(
57                    'Cannot determine board for host %s' % dut_host.hostname)
58        logging.debug('Detected board %s for host %s', board, dut_host.hostname)
59        return board
60
61    @staticmethod
62    def _construct_image_label(dut_board, release_version):
63        """Constructs a label combining the board name and release version.
64
65        @param dut_board: A string representing the board of the remote host.
66        @param release_version: A chromeOS release version.
67
68        @return: A string representing the release version.
69                 Ex: lumpy-release/R28-3993.0.0
70        """
71        # todo(rpius): We should probably make this more flexible to accept
72        # images from trybot's, etc.
73        return dut_board + '-release/' + release_version
74
75    @staticmethod
76    def _get_update_url(ds_url, image):
77        """Returns the full update URL. """
78        config = global_config.global_config
79        image_url_pattern = config.get_config_value(
80                'CROS', 'image_url_pattern', type=str)
81        return image_url_pattern % (ds_url, image)
82
83    @staticmethod
84    def _get_release_version_from_dut(dut_host):
85        """Get release version from the DUT located in lsb-release file.
86
87        @param dut_host: Host object representing the DUT.
88
89        @return: A string representing the release version.
90        """
91        return dut_host.get_release_version()
92
93    @staticmethod
94    def _get_release_version_from_image(image):
95        """Get release version from the image label.
96
97        @param image: The build type and version to install on the host.
98
99        @return: A string representing the release version.
100        """
101        return image.split('-')[-1]
102
103    @staticmethod
104    def _get_latest_release_version_from_server(dut_board):
105        """Gets the latest release version for a given board from a dev server.
106
107        @param dut_board: A string representing the board of the remote host.
108
109        @return: A string representing the release version.
110        """
111        build_target = dut_board + "-release"
112        config = global_config.global_config
113        server_url_list = config.get_config_value(
114                'CROS', 'dev_server', type=list, default=[])
115        ds = dev_server.ImageServer(server_url_list[0])
116        return ds.get_latest_build_in_server(build_target)
117
118    def update_dut(self, dut_host, image, force=True):
119        """The method called by to start the upgrade of a single DUT.
120
121        @param dut_host: Host object representing the DUT.
122        @param image: The build type and version to install on the host.
123        @param force: If False, will only updated the host if it is not
124                      already running the build. If True, force the
125                      update regardless, and force a full-reimage.
126
127        """
128        logging.debug('Host: %s. Start updating DUT to %s', dut_host, image)
129
130        # If the host is already on the correct build, we have nothing to do.
131        dut_release_version = self._get_release_version_from_dut(dut_host)
132        image_release_version = self._get_release_version_from_image(image)
133        if not force and dut_release_version == image_release_version:
134            logging.info('Host: %s. Already running %s',
135                         dut_host, image_release_version)
136            sys.exit(UPDATE_SUCCESS)
137
138        try:
139            ds = dev_server.ImageServer.resolve(image)
140            # We need the autotest packages to run the tests.
141            ds.stage_artifacts(image, ['full_payload', 'stateful',
142                                       'autotest_packages'])
143        except dev_server.DevServerException as e:
144            error_str = 'Host: ' + dut_host + '. ' + e
145            logging.error(error_str)
146            sys.exit(UPDATE_FAILURE)
147
148        url = self._get_update_url(ds.url(), image)
149        logging.debug('Host: %s. Installing image from %s', dut_host, url)
150        try:
151            autoupdater.ChromiumOSUpdater(url, host=dut_host).run_update()
152        except error.TestFail as e:
153            error_str = 'Host: ' + dut_host + '. ' + e
154            logging.error(error_str)
155            sys.exit(UPDATE_FAILURE)
156
157        dut_release_version = self._get_release_version_from_dut(dut_host)
158        if dut_release_version != image_release_version:
159            error_str = 'Host: ' + dut_host + '. Expected version of ' + \
160                        image_release_version + ' in DUT, but found '  + \
161                        dut_release_version + '.'
162            logging.error(error_str)
163            sys.exit(UPDATE_FAILURE)
164
165        logging.info('Host: %s. Finished updating DUT to %s', dut_host, image)
166        sys.exit(UPDATE_SUCCESS)
167
168    def update_dut_pool(self, dut_objects, release_version=""):
169        """Updates all the DUT's in the pool to a provided release version.
170
171        @param dut_objects: An array of DUTObjects corresponding to all the
172                            DUT's in the DUT pool.
173        @param release_version: A chromeOS release version.
174
175        @return: True if all the DUT's successfully upgraded, False otherwise.
176        """
177        tasks = []
178        for dut in dut_objects:
179            dut_board = self._get_board_name_from_host(dut.host)
180            if release_version == "":
181                release_version = self._get_latest_release_version_from_server(
182                        dut_board)
183            dut_image = self._construct_image_label(dut_board, release_version)
184            # Schedule the update for this DUT to the update process pool.
185            task = multiprocessing.Process(
186                    target=update_dut_worker,
187                    args=(self, dut, dut_image, False))
188            tasks.append(task)
189        # Run the updates in parallel.
190        for task in tasks:
191            task.start()
192        for task in tasks:
193            task.join()
194
195        # Check the exit code to determine if the updates were all successful
196        # or not.
197        for task in tasks:
198            if task.exitcode == UPDATE_FAILURE:
199                return False
200        return True
201