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.dynamic_suite import constants
13
14#Update status
15UPDATE_SUCCESS = 0
16UPDATE_FAILURE = 1
17
18def update_dut_worker(updater_obj, dut, image, force):
19    """The method called by multiprocessing worker pool for updating DUT.
20    This function is the function which is repeatedly scheduled for each
21    DUT through the multiprocessing worker. This has to be defined outside
22    the class because it needs to be pickleable.
23
24    @param updater_obj: An CliqueDUTUpdater object.
25    @param dut: DUTObject representing the DUT.
26    @param image: The build type and version to install on the host.
27    @param force: If False, will only updated the host if it is not
28                  already running the build. If True, force the
29                  update regardless, and force a full-reimage.
30
31    """
32    updater_obj.update_dut(dut_host=dut.host, image=image, force=force)
33
34
35class CliqueDUTUpdater(object):
36    """CliqueDUTUpdater is responsible for updating all the DUT's in the
37    DUT pool to the same release.
38    """
39
40    def __init__(self):
41        """Initializes the DUT updater for updating the DUT's in the pool."""
42
43
44    @staticmethod
45    def _get_board_name_from_host(dut_host):
46        """Get the board name of the remote host.
47
48        @param host: Host object representing the DUT.
49
50        @return: A string representing the board of the remote host.
51        """
52        try:
53            board = dut_host.get_board().replace(constants.BOARD_PREFIX, '')
54        except error.AutoservRunError:
55            raise error.TestFail(
56                    'Cannot determine board for host %s' % dut_host.hostname)
57        logging.debug('Detected board %s for host %s', board, dut_host.hostname)
58        return board
59
60    @staticmethod
61    def _construct_image_label(dut_board, release_version):
62        """Constructs a label combining the board name and release version.
63
64        @param dut_board: A string representing the board of the remote host.
65        @param release_version: A chromeOS release version.
66
67        @return: A string representing the release version.
68                 Ex: lumpy-release/R28-3993.0.0
69        """
70        # todo(rpius): We should probably make this more flexible to accept
71        # images from trybot's, etc.
72        return dut_board + '-release/' + release_version
73
74    @staticmethod
75    def _get_update_url(ds_url, image):
76        """Returns the full update URL. """
77        config = global_config.global_config
78        image_url_pattern = config.get_config_value(
79                'CROS', 'image_url_pattern', type=str)
80        return image_url_pattern % (ds_url, image)
81
82    @staticmethod
83    def _get_release_version_from_dut(dut_host):
84        """Get release version from the DUT located in lsb-release file.
85
86        @param dut_host: Host object representing the DUT.
87
88        @return: A string representing the release version.
89        """
90        return dut_host.get_release_version()
91
92    @staticmethod
93    def _get_release_version_from_image(image):
94        """Get release version from the image label.
95
96        @param image: The build type and version to install on the host.
97
98        @return: A string representing the release version.
99        """
100        return image.split('-')[-1]
101
102    @staticmethod
103    def _get_latest_release_version_from_server(dut_board):
104        """Gets the latest release version for a given board from a dev server.
105
106        @param dut_board: A string representing the board of the remote host.
107
108        @return: A string representing the release version.
109        """
110        build_target = dut_board + "-release"
111        config = global_config.global_config
112        server_url_list = config.get_config_value(
113                'CROS', 'dev_server', type=list, default=[])
114        ds = dev_server.ImageServer(server_url_list[0])
115        return ds.get_latest_build_in_server(build_target)
116
117    def update_dut(self, dut_host, image, force=True):
118        """The method called by to start the upgrade of a single DUT.
119
120        @param dut_host: Host object representing the DUT.
121        @param image: The build type and version to install on the host.
122        @param force: If False, will only updated the host if it is not
123                      already running the build. If True, force the
124                      update regardless, and force a full-reimage.
125
126        """
127        logging.debug('Host: %s. Start updating DUT to %s', dut_host, image)
128
129        # If the host is already on the correct build, we have nothing to do.
130        dut_release_version = self._get_release_version_from_dut(dut_host)
131        image_release_version = self._get_release_version_from_image(image)
132        if not force and dut_release_version == image_release_version:
133            logging.info('Host: %s. Already running %s',
134                         dut_host, image_release_version)
135            sys.exit(UPDATE_SUCCESS)
136
137        try:
138            ds = dev_server.ImageServer.resolve(image)
139            # We need the autotest packages to run the tests.
140            ds.stage_artifacts(image, ['full_payload', 'stateful',
141                                       'autotest_packages'])
142        except dev_server.DevServerException as e:
143            error_str = 'Host: ' + dut_host + '. ' + e
144            logging.error(error_str)
145            sys.exit(UPDATE_FAILURE)
146
147        url = self._get_update_url(ds.url(), image)
148        logging.debug('Host: %s. Installing image from %s', dut_host, url)
149        try:
150            dut_host.machine_install(force_update=True, update_url=url,
151                                     force_full_update=force)
152        except error.InstallError 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