xref: /aosp_15_r20/external/autotest/server/frontend.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1*9c5db199SXin Li# Lint as: python2, python3
2*9c5db199SXin Li# Copyright Martin J. Bligh, Google Inc 2008
3*9c5db199SXin Li# Released under the GPL v2
4*9c5db199SXin Li
5*9c5db199SXin Li"""
6*9c5db199SXin LiThis class allows you to communicate with the frontend to submit jobs etc
7*9c5db199SXin LiIt is designed for writing more sophisiticated server-side control files that
8*9c5db199SXin Lican recursively add and manage other jobs.
9*9c5db199SXin Li
10*9c5db199SXin LiWe turn the JSON dictionaries into real objects that are more idiomatic
11*9c5db199SXin Li
12*9c5db199SXin LiFor docs, see:
13*9c5db199SXin Li    http://www.chromium.org/chromium-os/testing/afe-rpc-infrastructure
14*9c5db199SXin Li    http://docs.djangoproject.com/en/dev/ref/models/querysets/#queryset-api
15*9c5db199SXin Li"""
16*9c5db199SXin Li
17*9c5db199SXin Li#pylint: disable=missing-docstring
18*9c5db199SXin Li
19*9c5db199SXin Lifrom __future__ import absolute_import
20*9c5db199SXin Lifrom __future__ import division
21*9c5db199SXin Lifrom __future__ import print_function
22*9c5db199SXin Li
23*9c5db199SXin Liimport getpass
24*9c5db199SXin Liimport os
25*9c5db199SXin Liimport re
26*9c5db199SXin Li
27*9c5db199SXin Liimport common
28*9c5db199SXin Li
29*9c5db199SXin Lifrom autotest_lib.frontend.afe import rpc_client_lib
30*9c5db199SXin Lifrom autotest_lib.client.common_lib import control_data
31*9c5db199SXin Lifrom autotest_lib.client.common_lib import global_config
32*9c5db199SXin Lifrom autotest_lib.client.common_lib import host_states
33*9c5db199SXin Lifrom autotest_lib.client.common_lib import priorities
34*9c5db199SXin Lifrom autotest_lib.client.common_lib import utils
35*9c5db199SXin Lifrom autotest_lib.tko import db
36*9c5db199SXin Lifrom six.moves import zip
37*9c5db199SXin Li
38*9c5db199SXin Litry:
39*9c5db199SXin Li    from autotest_lib.utils.frozen_chromite.lib import metrics
40*9c5db199SXin Liexcept ImportError:
41*9c5db199SXin Li    metrics = utils.metrics_mock
42*9c5db199SXin Li
43*9c5db199SXin Litry:
44*9c5db199SXin Li    from autotest_lib.server.site_common import site_utils as server_utils
45*9c5db199SXin Liexcept:
46*9c5db199SXin Li    from autotest_lib.server import utils as server_utils
47*9c5db199SXin Liform_ntuples_from_machines = server_utils.form_ntuples_from_machines
48*9c5db199SXin Li
49*9c5db199SXin LiGLOBAL_CONFIG = global_config.global_config
50*9c5db199SXin LiDEFAULT_SERVER = 'autotest'
51*9c5db199SXin Li
52*9c5db199SXin Li
53*9c5db199SXin Lidef dump_object(header, obj):
54*9c5db199SXin Li    """
55*9c5db199SXin Li    Standard way to print out the frontend objects (eg job, host, acl, label)
56*9c5db199SXin Li    in a human-readable fashion for debugging
57*9c5db199SXin Li    """
58*9c5db199SXin Li    result = header + '\n'
59*9c5db199SXin Li    for key in obj.hash:
60*9c5db199SXin Li        if key == 'afe' or key == 'hash':
61*9c5db199SXin Li            continue
62*9c5db199SXin Li        result += '%20s: %s\n' % (key, obj.hash[key])
63*9c5db199SXin Li    return result
64*9c5db199SXin Li
65*9c5db199SXin Li
66*9c5db199SXin Liclass RpcClient(object):
67*9c5db199SXin Li    """
68*9c5db199SXin Li    Abstract RPC class for communicating with the autotest frontend
69*9c5db199SXin Li    Inherited for both TKO and AFE uses.
70*9c5db199SXin Li
71*9c5db199SXin Li    All the constructors go in the afe / tko class.
72*9c5db199SXin Li    Manipulating methods go in the object classes themselves
73*9c5db199SXin Li    """
74*9c5db199SXin Li    def __init__(self, path, user, server, print_log, debug, reply_debug):
75*9c5db199SXin Li        """
76*9c5db199SXin Li        Create a cached instance of a connection to the frontend
77*9c5db199SXin Li
78*9c5db199SXin Li            user: username to connect as
79*9c5db199SXin Li            server: frontend server to connect to
80*9c5db199SXin Li            print_log: pring a logging message to stdout on every operation
81*9c5db199SXin Li            debug: print out all RPC traffic
82*9c5db199SXin Li        """
83*9c5db199SXin Li        if not user and utils.is_in_container():
84*9c5db199SXin Li            user = GLOBAL_CONFIG.get_config_value('SSP', 'user', default=None)
85*9c5db199SXin Li        if not user:
86*9c5db199SXin Li            user = getpass.getuser()
87*9c5db199SXin Li        if not server:
88*9c5db199SXin Li            if 'AUTOTEST_WEB' in os.environ:
89*9c5db199SXin Li                server = os.environ['AUTOTEST_WEB']
90*9c5db199SXin Li            else:
91*9c5db199SXin Li                server = GLOBAL_CONFIG.get_config_value('SERVER', 'hostname',
92*9c5db199SXin Li                                                        default=DEFAULT_SERVER)
93*9c5db199SXin Li        self.server = server
94*9c5db199SXin Li        self.user = user
95*9c5db199SXin Li        self.print_log = print_log
96*9c5db199SXin Li        self.debug = debug
97*9c5db199SXin Li        self.reply_debug = reply_debug
98*9c5db199SXin Li        headers = {'AUTHORIZATION': self.user}
99*9c5db199SXin Li        rpc_server = rpc_client_lib.add_protocol(server) + path
100*9c5db199SXin Li        if debug:
101*9c5db199SXin Li            print('SERVER: %s' % rpc_server)
102*9c5db199SXin Li            print('HEADERS: %s' % headers)
103*9c5db199SXin Li        self.proxy = rpc_client_lib.get_proxy(rpc_server, headers=headers)
104*9c5db199SXin Li
105*9c5db199SXin Li
106*9c5db199SXin Li    def run(self, call, **dargs):
107*9c5db199SXin Li        """
108*9c5db199SXin Li        Make a RPC call to the AFE server
109*9c5db199SXin Li        """
110*9c5db199SXin Li        rpc_call = getattr(self.proxy, call)
111*9c5db199SXin Li        if self.debug:
112*9c5db199SXin Li            print('DEBUG: %s %s' % (call, dargs))
113*9c5db199SXin Li        try:
114*9c5db199SXin Li            result = utils.strip_unicode(rpc_call(**dargs))
115*9c5db199SXin Li            if self.reply_debug:
116*9c5db199SXin Li                print(result)
117*9c5db199SXin Li            return result
118*9c5db199SXin Li        except Exception:
119*9c5db199SXin Li            raise
120*9c5db199SXin Li
121*9c5db199SXin Li
122*9c5db199SXin Li    def log(self, message):
123*9c5db199SXin Li        if self.print_log:
124*9c5db199SXin Li            print(message)
125*9c5db199SXin Li
126*9c5db199SXin Li
127*9c5db199SXin Liclass TKO(RpcClient):
128*9c5db199SXin Li    def __init__(self, user=None, server=None, print_log=True, debug=False,
129*9c5db199SXin Li                 reply_debug=False):
130*9c5db199SXin Li        super(TKO, self).__init__(path='/new_tko/server/noauth/rpc/',
131*9c5db199SXin Li                                  user=user,
132*9c5db199SXin Li                                  server=server,
133*9c5db199SXin Li                                  print_log=print_log,
134*9c5db199SXin Li                                  debug=debug,
135*9c5db199SXin Li                                  reply_debug=reply_debug)
136*9c5db199SXin Li        self._db = None
137*9c5db199SXin Li
138*9c5db199SXin Li
139*9c5db199SXin Li    @metrics.SecondsTimerDecorator(
140*9c5db199SXin Li            'chromeos/autotest/tko/get_job_status_duration')
141*9c5db199SXin Li    def get_job_test_statuses_from_db(self, job_id):
142*9c5db199SXin Li        """Get job test statuses from the database.
143*9c5db199SXin Li
144*9c5db199SXin Li        Retrieve a set of fields from a job that reflect the status of each test
145*9c5db199SXin Li        run within a job.
146*9c5db199SXin Li        fields retrieved: status, test_name, reason, test_started_time,
147*9c5db199SXin Li                          test_finished_time, afe_job_id, job_owner, hostname.
148*9c5db199SXin Li
149*9c5db199SXin Li        @param job_id: The afe job id to look up.
150*9c5db199SXin Li        @returns a TestStatus object of the resulting information.
151*9c5db199SXin Li        """
152*9c5db199SXin Li        if self._db is None:
153*9c5db199SXin Li            self._db = db.db()
154*9c5db199SXin Li        fields = ['status', 'test_name', 'subdir', 'reason',
155*9c5db199SXin Li                  'test_started_time', 'test_finished_time', 'afe_job_id',
156*9c5db199SXin Li                  'job_owner', 'hostname', 'job_tag']
157*9c5db199SXin Li        table = 'tko_test_view_2'
158*9c5db199SXin Li        where = 'job_tag like "%s-%%"' % job_id
159*9c5db199SXin Li        test_status = []
160*9c5db199SXin Li        # Run commit before we query to ensure that we are pulling the latest
161*9c5db199SXin Li        # results.
162*9c5db199SXin Li        self._db.commit()
163*9c5db199SXin Li        for entry in self._db.select(','.join(fields), table, (where, None)):
164*9c5db199SXin Li            status_dict = {}
165*9c5db199SXin Li            for key,value in zip(fields, entry):
166*9c5db199SXin Li                # All callers expect values to be a str object.
167*9c5db199SXin Li                status_dict[key] = str(value)
168*9c5db199SXin Li            # id is used by TestStatus to uniquely identify each Test Status
169*9c5db199SXin Li            # obj.
170*9c5db199SXin Li            status_dict['id'] = [status_dict['reason'], status_dict['hostname'],
171*9c5db199SXin Li                                 status_dict['test_name']]
172*9c5db199SXin Li            test_status.append(status_dict)
173*9c5db199SXin Li
174*9c5db199SXin Li        return [TestStatus(self, e) for e in test_status]
175*9c5db199SXin Li
176*9c5db199SXin Li
177*9c5db199SXin Li    def get_status_counts(self, job, **data):
178*9c5db199SXin Li        entries = self.run('get_status_counts',
179*9c5db199SXin Li                           group_by=['hostname', 'test_name', 'reason'],
180*9c5db199SXin Li                           job_tag__startswith='%s-' % job, **data)
181*9c5db199SXin Li        return [TestStatus(self, e) for e in entries['groups']]
182*9c5db199SXin Li
183*9c5db199SXin Li
184*9c5db199SXin Liclass _StableVersionMap(object):
185*9c5db199SXin Li    """
186*9c5db199SXin Li    A mapping from board names to strings naming software versions.
187*9c5db199SXin Li
188*9c5db199SXin Li    The mapping is meant to allow finding a nominally "stable" version
189*9c5db199SXin Li    of software associated with a given board.  The mapping identifies
190*9c5db199SXin Li    specific versions of software that should be installed during
191*9c5db199SXin Li    operations such as repair.
192*9c5db199SXin Li
193*9c5db199SXin Li    Conceptually, there are multiple version maps, each handling
194*9c5db199SXin Li    different types of image.  For instance, a single board may have
195*9c5db199SXin Li    both a stable OS image (e.g. for CrOS), and a separate stable
196*9c5db199SXin Li    firmware image.
197*9c5db199SXin Li
198*9c5db199SXin Li    Each different type of image requires a certain amount of special
199*9c5db199SXin Li    handling, implemented by a subclass of `StableVersionMap`.  The
200*9c5db199SXin Li    subclasses take care of pre-processing of arguments, delegating
201*9c5db199SXin Li    actual RPC calls to this superclass.
202*9c5db199SXin Li
203*9c5db199SXin Li    @property _afe      AFE object through which to make the actual RPC
204*9c5db199SXin Li                        calls.
205*9c5db199SXin Li    @property _android  Value of the `android` parameter to be passed
206*9c5db199SXin Li                        when calling the `get_stable_version` RPC.
207*9c5db199SXin Li    """
208*9c5db199SXin Li
209*9c5db199SXin Li    def __init__(self, afe):
210*9c5db199SXin Li        self._afe = afe
211*9c5db199SXin Li
212*9c5db199SXin Li
213*9c5db199SXin Li    def get_all_versions(self):
214*9c5db199SXin Li        """
215*9c5db199SXin Li        Get all mappings in the stable versions table.
216*9c5db199SXin Li
217*9c5db199SXin Li        Extracts the full content of the `stable_version` table
218*9c5db199SXin Li        in the AFE database, and returns it as a dictionary
219*9c5db199SXin Li        mapping board names to version strings.
220*9c5db199SXin Li
221*9c5db199SXin Li        @return A dictionary mapping board names to version strings.
222*9c5db199SXin Li        """
223*9c5db199SXin Li        return self._afe.run('get_all_stable_versions')
224*9c5db199SXin Li
225*9c5db199SXin Li
226*9c5db199SXin Li    def get_version(self, board):
227*9c5db199SXin Li        """
228*9c5db199SXin Li        Get the mapping of one board in the stable versions table.
229*9c5db199SXin Li
230*9c5db199SXin Li        Look up and return the version mapped to the given board in the
231*9c5db199SXin Li        `stable_versions` table in the AFE database.
232*9c5db199SXin Li
233*9c5db199SXin Li        @param board  The board to be looked up.
234*9c5db199SXin Li
235*9c5db199SXin Li        @return The version mapped for the given board.
236*9c5db199SXin Li        """
237*9c5db199SXin Li        return self._afe.run('get_stable_version', board=board)
238*9c5db199SXin Li
239*9c5db199SXin Li
240*9c5db199SXin Li    def set_version(self, board, version):
241*9c5db199SXin Li        """
242*9c5db199SXin Li        Change the mapping of one board in the stable versions table.
243*9c5db199SXin Li
244*9c5db199SXin Li        Set the mapping in the `stable_versions` table in the AFE
245*9c5db199SXin Li        database for the given board to the given version.
246*9c5db199SXin Li
247*9c5db199SXin Li        @param board    The board to be updated.
248*9c5db199SXin Li        @param version  The new version to be assigned to the board.
249*9c5db199SXin Li        """
250*9c5db199SXin Li        raise RuntimeError("server.frontend._StableVersionMap::set_version is intentionally deleted")
251*9c5db199SXin Li
252*9c5db199SXin Li
253*9c5db199SXin Li    def delete_version(self, board):
254*9c5db199SXin Li        """
255*9c5db199SXin Li        Remove the mapping of one board in the stable versions table.
256*9c5db199SXin Li
257*9c5db199SXin Li        Remove the mapping in the `stable_versions` table in the AFE
258*9c5db199SXin Li        database for the given board.
259*9c5db199SXin Li
260*9c5db199SXin Li        @param board    The board to be updated.
261*9c5db199SXin Li        """
262*9c5db199SXin Li        raise RuntimeError("server.frontend._StableVersionMap::delete_version is intentionally deleted")
263*9c5db199SXin Li
264*9c5db199SXin Li
265*9c5db199SXin Liclass _OSVersionMap(_StableVersionMap):
266*9c5db199SXin Li    """
267*9c5db199SXin Li    Abstract stable version mapping for full OS images of various types.
268*9c5db199SXin Li    """
269*9c5db199SXin Li
270*9c5db199SXin Li    def _version_is_valid(self, version):
271*9c5db199SXin Li        return True
272*9c5db199SXin Li
273*9c5db199SXin Li    def get_all_versions(self):
274*9c5db199SXin Li        versions = super(_OSVersionMap, self).get_all_versions()
275*9c5db199SXin Li        for board in versions.keys():
276*9c5db199SXin Li            if ('/' in board
277*9c5db199SXin Li                    or not self._version_is_valid(versions[board])):
278*9c5db199SXin Li                del versions[board]
279*9c5db199SXin Li        return versions
280*9c5db199SXin Li
281*9c5db199SXin Li    def get_version(self, board):
282*9c5db199SXin Li        version = super(_OSVersionMap, self).get_version(board)
283*9c5db199SXin Li        return version if self._version_is_valid(version) else None
284*9c5db199SXin Li
285*9c5db199SXin Li
286*9c5db199SXin Lidef format_cros_image_name(board, version):
287*9c5db199SXin Li    """
288*9c5db199SXin Li    Return an image name for a given `board` and `version`.
289*9c5db199SXin Li
290*9c5db199SXin Li    This formats `board` and `version` into a string identifying an
291*9c5db199SXin Li    image file.  The string represents part of a URL for access to
292*9c5db199SXin Li    the image.
293*9c5db199SXin Li
294*9c5db199SXin Li    The returned image name is typically of a form like
295*9c5db199SXin Li    "falco-release/R55-8872.44.0".
296*9c5db199SXin Li    """
297*9c5db199SXin Li    build_pattern = GLOBAL_CONFIG.get_config_value(
298*9c5db199SXin Li            'CROS', 'stable_build_pattern')
299*9c5db199SXin Li    return build_pattern % (board, version)
300*9c5db199SXin Li
301*9c5db199SXin Li
302*9c5db199SXin Liclass _CrosVersionMap(_OSVersionMap):
303*9c5db199SXin Li    """
304*9c5db199SXin Li    Stable version mapping for ChromeOS release images.
305*9c5db199SXin Li
306*9c5db199SXin Li    This class manages a mapping of ChromeOS board names to known-good
307*9c5db199SXin Li    release (or canary) images.  The images selected can be installed on
308*9c5db199SXin Li    DUTs during repair tasks, as a way of getting a DUT into a known
309*9c5db199SXin Li    working state.
310*9c5db199SXin Li    """
311*9c5db199SXin Li
312*9c5db199SXin Li    def _version_is_valid(self, version):
313*9c5db199SXin Li        return version is not None and '/' not in version
314*9c5db199SXin Li
315*9c5db199SXin Li    def get_image_name(self, board):
316*9c5db199SXin Li        """
317*9c5db199SXin Li        Return the full image name of the stable version for `board`.
318*9c5db199SXin Li
319*9c5db199SXin Li        This finds the stable version for `board`, and returns a string
320*9c5db199SXin Li        identifying the associated image as for `format_image_name()`,
321*9c5db199SXin Li        above.
322*9c5db199SXin Li
323*9c5db199SXin Li        @return A string identifying the image file for the stable
324*9c5db199SXin Li                image for `board`.
325*9c5db199SXin Li        """
326*9c5db199SXin Li        return format_cros_image_name(board, self.get_version(board))
327*9c5db199SXin Li
328*9c5db199SXin Li
329*9c5db199SXin Liclass _SuffixHackVersionMap(_StableVersionMap):
330*9c5db199SXin Li    """
331*9c5db199SXin Li    Abstract super class for mappings using a pseudo-board name.
332*9c5db199SXin Li
333*9c5db199SXin Li    For non-OS image type mappings, we look them up in the
334*9c5db199SXin Li    `stable_versions` table by constructing a "pseudo-board" from the
335*9c5db199SXin Li    real board name plus a suffix string that identifies the image type.
336*9c5db199SXin Li    So, for instance the name "lulu/firmware" is used to look up the
337*9c5db199SXin Li    FAFT firmware version for lulu boards.
338*9c5db199SXin Li    """
339*9c5db199SXin Li
340*9c5db199SXin Li    # _SUFFIX - The suffix used in constructing the "pseudo-board"
341*9c5db199SXin Li    # lookup key.  Each subclass must define this value for itself.
342*9c5db199SXin Li    #
343*9c5db199SXin Li    _SUFFIX = None
344*9c5db199SXin Li
345*9c5db199SXin Li    def get_all_versions(self):
346*9c5db199SXin Li        # Get all the mappings from the AFE, extract just the mappings
347*9c5db199SXin Li        # with our suffix, and replace the pseudo-board name keys with
348*9c5db199SXin Li        # the real board names.
349*9c5db199SXin Li        #
350*9c5db199SXin Li        all_versions = super(
351*9c5db199SXin Li                _SuffixHackVersionMap, self).get_all_versions()
352*9c5db199SXin Li        return {
353*9c5db199SXin Li            board[0 : -len(self._SUFFIX)]: all_versions[board]
354*9c5db199SXin Li                for board in all_versions.keys()
355*9c5db199SXin Li                    if board.endswith(self._SUFFIX)
356*9c5db199SXin Li        }
357*9c5db199SXin Li
358*9c5db199SXin Li
359*9c5db199SXin Li    def get_version(self, board):
360*9c5db199SXin Li        board += self._SUFFIX
361*9c5db199SXin Li        return super(_SuffixHackVersionMap, self).get_version(board)
362*9c5db199SXin Li
363*9c5db199SXin Li
364*9c5db199SXin Li    def set_version(self, board, version):
365*9c5db199SXin Li        board += self._SUFFIX
366*9c5db199SXin Li        super(_SuffixHackVersionMap, self).set_version(board, version)
367*9c5db199SXin Li
368*9c5db199SXin Li
369*9c5db199SXin Li    def delete_version(self, board):
370*9c5db199SXin Li        board += self._SUFFIX
371*9c5db199SXin Li        super(_SuffixHackVersionMap, self).delete_version(board)
372*9c5db199SXin Li
373*9c5db199SXin Li
374*9c5db199SXin Liclass _FAFTVersionMap(_SuffixHackVersionMap):
375*9c5db199SXin Li    """
376*9c5db199SXin Li    Stable version mapping for firmware versions used in FAFT repair.
377*9c5db199SXin Li
378*9c5db199SXin Li    When DUTs used for FAFT fail repair, stable firmware may need to be
379*9c5db199SXin Li    flashed directly from original tarballs.  The FAFT firmware version
380*9c5db199SXin Li    mapping finds the appropriate tarball for a given board.
381*9c5db199SXin Li    """
382*9c5db199SXin Li
383*9c5db199SXin Li    _SUFFIX = '/firmware'
384*9c5db199SXin Li
385*9c5db199SXin Li    def get_version(self, board):
386*9c5db199SXin Li        # If there's no mapping for `board`, the lookup will return the
387*9c5db199SXin Li        # default CrOS version mapping.  To eliminate that case, we
388*9c5db199SXin Li        # require a '/' character in the version, since CrOS versions
389*9c5db199SXin Li        # won't match that.
390*9c5db199SXin Li        #
391*9c5db199SXin Li        # TODO(jrbarnette):  This is, of course, a hack.  Ultimately,
392*9c5db199SXin Li        # the right fix is to move handling to the RPC server side.
393*9c5db199SXin Li        #
394*9c5db199SXin Li        version = super(_FAFTVersionMap, self).get_version(board)
395*9c5db199SXin Li        return version if '/' in version else None
396*9c5db199SXin Li
397*9c5db199SXin Li
398*9c5db199SXin Liclass _FirmwareVersionMap(_SuffixHackVersionMap):
399*9c5db199SXin Li    """
400*9c5db199SXin Li    Stable version mapping for firmware supplied in ChromeOS images.
401*9c5db199SXin Li
402*9c5db199SXin Li    A ChromeOS image bundles a version of the firmware that the
403*9c5db199SXin Li    device should update to when the OS version is installed during
404*9c5db199SXin Li    AU.
405*9c5db199SXin Li
406*9c5db199SXin Li    Test images suppress the firmware update during AU.  Instead, during
407*9c5db199SXin Li    repair and verify we check installed firmware on a DUT, compare it
408*9c5db199SXin Li    against the stable version mapping for the board, and update when
409*9c5db199SXin Li    the DUT is out-of-date.
410*9c5db199SXin Li    """
411*9c5db199SXin Li
412*9c5db199SXin Li    _SUFFIX = '/rwfw'
413*9c5db199SXin Li
414*9c5db199SXin Li    def get_version(self, board):
415*9c5db199SXin Li        # If there's no mapping for `board`, the lookup will return the
416*9c5db199SXin Li        # default CrOS version mapping.  To eliminate that case, we
417*9c5db199SXin Li        # require the version start with "Google_", since CrOS versions
418*9c5db199SXin Li        # won't match that.
419*9c5db199SXin Li        #
420*9c5db199SXin Li        # TODO(jrbarnette):  This is, of course, a hack.  Ultimately,
421*9c5db199SXin Li        # the right fix is to move handling to the RPC server side.
422*9c5db199SXin Li        #
423*9c5db199SXin Li        version = super(_FirmwareVersionMap, self).get_version(board)
424*9c5db199SXin Li        return version if version.startswith('Google_') else None
425*9c5db199SXin Li
426*9c5db199SXin Li
427*9c5db199SXin Liclass AFE(RpcClient):
428*9c5db199SXin Li
429*9c5db199SXin Li    # Known image types for stable version mapping objects.
430*9c5db199SXin Li    # CROS_IMAGE_TYPE - Mappings for ChromeOS images.
431*9c5db199SXin Li    # FAFT_IMAGE_TYPE - Mappings for Firmware images for FAFT repair.
432*9c5db199SXin Li    # FIRMWARE_IMAGE_TYPE - Mappings for released RW Firmware images.
433*9c5db199SXin Li    #
434*9c5db199SXin Li    CROS_IMAGE_TYPE = 'cros'
435*9c5db199SXin Li    FAFT_IMAGE_TYPE = 'faft'
436*9c5db199SXin Li    FIRMWARE_IMAGE_TYPE = 'firmware'
437*9c5db199SXin Li
438*9c5db199SXin Li    _IMAGE_MAPPING_CLASSES = {
439*9c5db199SXin Li        CROS_IMAGE_TYPE: _CrosVersionMap,
440*9c5db199SXin Li        FAFT_IMAGE_TYPE: _FAFTVersionMap,
441*9c5db199SXin Li        FIRMWARE_IMAGE_TYPE: _FirmwareVersionMap,
442*9c5db199SXin Li    }
443*9c5db199SXin Li
444*9c5db199SXin Li
445*9c5db199SXin Li    def __init__(self, user=None, server=None, print_log=True, debug=False,
446*9c5db199SXin Li                 reply_debug=False, job=None):
447*9c5db199SXin Li        self.job = job
448*9c5db199SXin Li        super(AFE, self).__init__(path='/afe/server/noauth/rpc/',
449*9c5db199SXin Li                                  user=user,
450*9c5db199SXin Li                                  server=server,
451*9c5db199SXin Li                                  print_log=print_log,
452*9c5db199SXin Li                                  debug=debug,
453*9c5db199SXin Li                                  reply_debug=reply_debug)
454*9c5db199SXin Li
455*9c5db199SXin Li
456*9c5db199SXin Li    def get_stable_version_map(self, image_type):
457*9c5db199SXin Li        """
458*9c5db199SXin Li        Return a stable version mapping for the given image type.
459*9c5db199SXin Li
460*9c5db199SXin Li        @return An object mapping board names to version strings for
461*9c5db199SXin Li                software of the given image type.
462*9c5db199SXin Li        """
463*9c5db199SXin Li        return self._IMAGE_MAPPING_CLASSES[image_type](self)
464*9c5db199SXin Li
465*9c5db199SXin Li
466*9c5db199SXin Li    def host_statuses(self, live=None):
467*9c5db199SXin Li        dead_statuses = ['Repair Failed', 'Repairing']
468*9c5db199SXin Li        statuses = self.run('get_static_data')['host_statuses']
469*9c5db199SXin Li        if live == True:
470*9c5db199SXin Li            return list(set(statuses) - set(dead_statuses))
471*9c5db199SXin Li        if live == False:
472*9c5db199SXin Li            return dead_statuses
473*9c5db199SXin Li        else:
474*9c5db199SXin Li            return statuses
475*9c5db199SXin Li
476*9c5db199SXin Li
477*9c5db199SXin Li    @staticmethod
478*9c5db199SXin Li    def _dict_for_host_query(hostnames=(), status=None, label=None):
479*9c5db199SXin Li        query_args = {}
480*9c5db199SXin Li        if hostnames:
481*9c5db199SXin Li            query_args['hostname__in'] = hostnames
482*9c5db199SXin Li        if status:
483*9c5db199SXin Li            query_args['status'] = status
484*9c5db199SXin Li        if label:
485*9c5db199SXin Li            query_args['labels__name'] = label
486*9c5db199SXin Li        return query_args
487*9c5db199SXin Li
488*9c5db199SXin Li
489*9c5db199SXin Li    def get_hosts(self, hostnames=(), status=None, label=None, **dargs):
490*9c5db199SXin Li        query_args = dict(dargs)
491*9c5db199SXin Li        query_args.update(self._dict_for_host_query(hostnames=hostnames,
492*9c5db199SXin Li                                                    status=status,
493*9c5db199SXin Li                                                    label=label))
494*9c5db199SXin Li        hosts = self.run('get_hosts', **query_args)
495*9c5db199SXin Li        return [Host(self, h) for h in hosts]
496*9c5db199SXin Li
497*9c5db199SXin Li
498*9c5db199SXin Li    def get_hostnames(self, status=None, label=None, **dargs):
499*9c5db199SXin Li        """Like get_hosts() but returns hostnames instead of Host objects."""
500*9c5db199SXin Li        # This implementation can be replaced with a more efficient one
501*9c5db199SXin Li        # that does not query for entire host objects in the future.
502*9c5db199SXin Li        return [host_obj.hostname for host_obj in
503*9c5db199SXin Li                self.get_hosts(status=status, label=label, **dargs)]
504*9c5db199SXin Li
505*9c5db199SXin Li
506*9c5db199SXin Li    def reverify_hosts(self, hostnames=(), status=None, label=None):
507*9c5db199SXin Li        query_args = dict(locked=False,
508*9c5db199SXin Li                          aclgroup__users__login=self.user)
509*9c5db199SXin Li        query_args.update(self._dict_for_host_query(hostnames=hostnames,
510*9c5db199SXin Li                                                    status=status,
511*9c5db199SXin Li                                                    label=label))
512*9c5db199SXin Li        return self.run('reverify_hosts', **query_args)
513*9c5db199SXin Li
514*9c5db199SXin Li
515*9c5db199SXin Li    def repair_hosts(self, hostnames=(), status=None, label=None):
516*9c5db199SXin Li        query_args = dict(locked=False,
517*9c5db199SXin Li                          aclgroup__users__login=self.user)
518*9c5db199SXin Li        query_args.update(self._dict_for_host_query(hostnames=hostnames,
519*9c5db199SXin Li                                                    status=status,
520*9c5db199SXin Li                                                    label=label))
521*9c5db199SXin Li        return self.run('repair_hosts', **query_args)
522*9c5db199SXin Li
523*9c5db199SXin Li
524*9c5db199SXin Li    def create_host(self, hostname, **dargs):
525*9c5db199SXin Li        id = self.run('add_host', hostname=hostname, **dargs)
526*9c5db199SXin Li        return self.get_hosts(id=id)[0]
527*9c5db199SXin Li
528*9c5db199SXin Li
529*9c5db199SXin Li    def get_host_attribute(self, attr, **dargs):
530*9c5db199SXin Li        host_attrs = self.run('get_host_attribute', attribute=attr, **dargs)
531*9c5db199SXin Li        return [HostAttribute(self, a) for a in host_attrs]
532*9c5db199SXin Li
533*9c5db199SXin Li
534*9c5db199SXin Li    def set_host_attribute(self, attr, val, **dargs):
535*9c5db199SXin Li        self.run('set_host_attribute', attribute=attr, value=val, **dargs)
536*9c5db199SXin Li
537*9c5db199SXin Li
538*9c5db199SXin Li    def get_labels(self, **dargs):
539*9c5db199SXin Li        labels = self.run('get_labels', **dargs)
540*9c5db199SXin Li        return [Label(self, l) for l in labels]
541*9c5db199SXin Li
542*9c5db199SXin Li
543*9c5db199SXin Li    def create_label(self, name, **dargs):
544*9c5db199SXin Li        id = self.run('add_label', name=name, **dargs)
545*9c5db199SXin Li        return self.get_labels(id=id)[0]
546*9c5db199SXin Li
547*9c5db199SXin Li
548*9c5db199SXin Li    def get_acls(self, **dargs):
549*9c5db199SXin Li        acls = self.run('get_acl_groups', **dargs)
550*9c5db199SXin Li        return [Acl(self, a) for a in acls]
551*9c5db199SXin Li
552*9c5db199SXin Li
553*9c5db199SXin Li    def create_acl(self, name, **dargs):
554*9c5db199SXin Li        id = self.run('add_acl_group', name=name, **dargs)
555*9c5db199SXin Li        return self.get_acls(id=id)[0]
556*9c5db199SXin Li
557*9c5db199SXin Li
558*9c5db199SXin Li    def get_users(self, **dargs):
559*9c5db199SXin Li        users = self.run('get_users', **dargs)
560*9c5db199SXin Li        return [User(self, u) for u in users]
561*9c5db199SXin Li
562*9c5db199SXin Li
563*9c5db199SXin Li    def generate_control_file(self, tests, **dargs):
564*9c5db199SXin Li        ret = self.run('generate_control_file', tests=tests, **dargs)
565*9c5db199SXin Li        return ControlFile(self, ret)
566*9c5db199SXin Li
567*9c5db199SXin Li
568*9c5db199SXin Li    def get_jobs(self, summary=False, **dargs):
569*9c5db199SXin Li        if summary:
570*9c5db199SXin Li            jobs_data = self.run('get_jobs_summary', **dargs)
571*9c5db199SXin Li        else:
572*9c5db199SXin Li            jobs_data = self.run('get_jobs', **dargs)
573*9c5db199SXin Li        jobs = []
574*9c5db199SXin Li        for j in jobs_data:
575*9c5db199SXin Li            job = Job(self, j)
576*9c5db199SXin Li            # Set up some extra information defaults
577*9c5db199SXin Li            job.testname = re.sub('\s.*', '', job.name) # arbitrary default
578*9c5db199SXin Li            job.platform_results = {}
579*9c5db199SXin Li            job.platform_reasons = {}
580*9c5db199SXin Li            jobs.append(job)
581*9c5db199SXin Li        return jobs
582*9c5db199SXin Li
583*9c5db199SXin Li
584*9c5db199SXin Li    def get_host_queue_entries(self, **kwargs):
585*9c5db199SXin Li        """Find JobStatus objects matching some constraints.
586*9c5db199SXin Li
587*9c5db199SXin Li        @param **kwargs: Arguments to pass to the RPC
588*9c5db199SXin Li        """
589*9c5db199SXin Li        entries = self.run('get_host_queue_entries', **kwargs)
590*9c5db199SXin Li        return self._entries_to_statuses(entries)
591*9c5db199SXin Li
592*9c5db199SXin Li
593*9c5db199SXin Li    def get_host_queue_entries_by_insert_time(self, **kwargs):
594*9c5db199SXin Li        """Like get_host_queue_entries, but using the insert index table.
595*9c5db199SXin Li
596*9c5db199SXin Li        @param **kwargs: Arguments to pass to the RPC
597*9c5db199SXin Li        """
598*9c5db199SXin Li        entries = self.run('get_host_queue_entries_by_insert_time', **kwargs)
599*9c5db199SXin Li        return self._entries_to_statuses(entries)
600*9c5db199SXin Li
601*9c5db199SXin Li
602*9c5db199SXin Li    def _entries_to_statuses(self, entries):
603*9c5db199SXin Li        """Converts HQEs to JobStatuses
604*9c5db199SXin Li
605*9c5db199SXin Li        Sadly, get_host_queue_entries doesn't return platforms, we have
606*9c5db199SXin Li        to get those back from an explicit get_hosts queury, then patch
607*9c5db199SXin Li        the new host objects back into the host list.
608*9c5db199SXin Li
609*9c5db199SXin Li        :param entries: A list of HQEs from get_host_queue_entries or
610*9c5db199SXin Li          get_host_queue_entries_by_insert_time.
611*9c5db199SXin Li        """
612*9c5db199SXin Li        job_statuses = [JobStatus(self, e) for e in entries]
613*9c5db199SXin Li        hostnames = [s.host.hostname for s in job_statuses if s.host]
614*9c5db199SXin Li        hosts = {}
615*9c5db199SXin Li        for host in self.get_hosts(hostname__in=hostnames):
616*9c5db199SXin Li            hosts[host.hostname] = host
617*9c5db199SXin Li        for status in job_statuses:
618*9c5db199SXin Li            if status.host:
619*9c5db199SXin Li                status.host = hosts.get(status.host.hostname)
620*9c5db199SXin Li        # filter job statuses that have either host or meta_host
621*9c5db199SXin Li        return [status for status in job_statuses if (status.host or
622*9c5db199SXin Li                                                      status.meta_host)]
623*9c5db199SXin Li
624*9c5db199SXin Li
625*9c5db199SXin Li    def get_special_tasks(self, **data):
626*9c5db199SXin Li        tasks = self.run('get_special_tasks', **data)
627*9c5db199SXin Li        return [SpecialTask(self, t) for t in tasks]
628*9c5db199SXin Li
629*9c5db199SXin Li
630*9c5db199SXin Li    def get_host_special_tasks(self, host_id, **data):
631*9c5db199SXin Li        tasks = self.run('get_host_special_tasks',
632*9c5db199SXin Li                         host_id=host_id, **data)
633*9c5db199SXin Li        return [SpecialTask(self, t) for t in tasks]
634*9c5db199SXin Li
635*9c5db199SXin Li
636*9c5db199SXin Li    def get_host_status_task(self, host_id, end_time):
637*9c5db199SXin Li        task = self.run('get_host_status_task',
638*9c5db199SXin Li                        host_id=host_id, end_time=end_time)
639*9c5db199SXin Li        return SpecialTask(self, task) if task else None
640*9c5db199SXin Li
641*9c5db199SXin Li
642*9c5db199SXin Li    def get_host_diagnosis_interval(self, host_id, end_time, success):
643*9c5db199SXin Li        return self.run('get_host_diagnosis_interval',
644*9c5db199SXin Li                        host_id=host_id, end_time=end_time,
645*9c5db199SXin Li                        success=success)
646*9c5db199SXin Li
647*9c5db199SXin Li
648*9c5db199SXin Li    def create_job(self, control_file, name=' ',
649*9c5db199SXin Li                   priority=priorities.Priority.DEFAULT,
650*9c5db199SXin Li                   control_type=control_data.CONTROL_TYPE_NAMES.CLIENT,
651*9c5db199SXin Li                   **dargs):
652*9c5db199SXin Li        id = self.run('create_job', name=name, priority=priority,
653*9c5db199SXin Li                 control_file=control_file, control_type=control_type, **dargs)
654*9c5db199SXin Li        return self.get_jobs(id=id)[0]
655*9c5db199SXin Li
656*9c5db199SXin Li
657*9c5db199SXin Li    def abort_jobs(self, jobs):
658*9c5db199SXin Li        """Abort a list of jobs.
659*9c5db199SXin Li
660*9c5db199SXin Li        Already completed jobs will not be affected.
661*9c5db199SXin Li
662*9c5db199SXin Li        @param jobs: List of job ids to abort.
663*9c5db199SXin Li        """
664*9c5db199SXin Li        for job in jobs:
665*9c5db199SXin Li            self.run('abort_host_queue_entries', job_id=job)
666*9c5db199SXin Li
667*9c5db199SXin Li
668*9c5db199SXin Li    def get_hosts_by_attribute(self, attribute, value):
669*9c5db199SXin Li        """
670*9c5db199SXin Li        Get the list of hosts that share the same host attribute value.
671*9c5db199SXin Li
672*9c5db199SXin Li        @param attribute: String of the host attribute to check.
673*9c5db199SXin Li        @param value: String of the value that is shared between hosts.
674*9c5db199SXin Li
675*9c5db199SXin Li        @returns List of hostnames that all have the same host attribute and
676*9c5db199SXin Li                 value.
677*9c5db199SXin Li        """
678*9c5db199SXin Li        return self.run('get_hosts_by_attribute',
679*9c5db199SXin Li                        attribute=attribute, value=value)
680*9c5db199SXin Li
681*9c5db199SXin Li
682*9c5db199SXin Li    def lock_host(self, host, lock_reason, fail_if_locked=False):
683*9c5db199SXin Li        """
684*9c5db199SXin Li        Lock the given host with the given lock reason.
685*9c5db199SXin Li
686*9c5db199SXin Li        Locking a host that's already locked using the 'modify_hosts' rpc
687*9c5db199SXin Li        will raise an exception. That's why fail_if_locked exists so the
688*9c5db199SXin Li        caller can determine if the lock succeeded or failed.  This will
689*9c5db199SXin Li        save every caller from wrapping lock_host in a try-except.
690*9c5db199SXin Li
691*9c5db199SXin Li        @param host: hostname of host to lock.
692*9c5db199SXin Li        @param lock_reason: Reason for locking host.
693*9c5db199SXin Li        @param fail_if_locked: Return False if host is already locked.
694*9c5db199SXin Li
695*9c5db199SXin Li        @returns Boolean, True if lock was successful, False otherwise.
696*9c5db199SXin Li        """
697*9c5db199SXin Li        try:
698*9c5db199SXin Li            self.run('modify_hosts',
699*9c5db199SXin Li                     host_filter_data={'hostname': host},
700*9c5db199SXin Li                     update_data={'locked': True,
701*9c5db199SXin Li                                  'lock_reason': lock_reason})
702*9c5db199SXin Li        except Exception:
703*9c5db199SXin Li            return not fail_if_locked
704*9c5db199SXin Li        return True
705*9c5db199SXin Li
706*9c5db199SXin Li
707*9c5db199SXin Li    def unlock_hosts(self, locked_hosts):
708*9c5db199SXin Li        """
709*9c5db199SXin Li        Unlock the hosts.
710*9c5db199SXin Li
711*9c5db199SXin Li        Unlocking a host that's already unlocked will do nothing so we don't
712*9c5db199SXin Li        need any special try-except clause here.
713*9c5db199SXin Li
714*9c5db199SXin Li        @param locked_hosts: List of hostnames of hosts to unlock.
715*9c5db199SXin Li        """
716*9c5db199SXin Li        self.run('modify_hosts',
717*9c5db199SXin Li                 host_filter_data={'hostname__in': locked_hosts},
718*9c5db199SXin Li                 update_data={'locked': False,
719*9c5db199SXin Li                              'lock_reason': ''})
720*9c5db199SXin Li
721*9c5db199SXin Li
722*9c5db199SXin Liclass TestResults(object):
723*9c5db199SXin Li    """
724*9c5db199SXin Li    Container class used to hold the results of the tests for a job
725*9c5db199SXin Li    """
726*9c5db199SXin Li    def __init__(self):
727*9c5db199SXin Li        self.good = []
728*9c5db199SXin Li        self.fail = []
729*9c5db199SXin Li        self.pending = []
730*9c5db199SXin Li
731*9c5db199SXin Li
732*9c5db199SXin Li    def add(self, result):
733*9c5db199SXin Li        if result.complete_count > result.pass_count:
734*9c5db199SXin Li            self.fail.append(result)
735*9c5db199SXin Li        elif result.incomplete_count > 0:
736*9c5db199SXin Li            self.pending.append(result)
737*9c5db199SXin Li        else:
738*9c5db199SXin Li            self.good.append(result)
739*9c5db199SXin Li
740*9c5db199SXin Li
741*9c5db199SXin Liclass RpcObject(object):
742*9c5db199SXin Li    """
743*9c5db199SXin Li    Generic object used to construct python objects from rpc calls
744*9c5db199SXin Li    """
745*9c5db199SXin Li    def __init__(self, afe, hash):
746*9c5db199SXin Li        self.afe = afe
747*9c5db199SXin Li        self.hash = hash
748*9c5db199SXin Li        self.__dict__.update(hash)
749*9c5db199SXin Li
750*9c5db199SXin Li
751*9c5db199SXin Li    def __str__(self):
752*9c5db199SXin Li        return dump_object(self.__repr__(), self)
753*9c5db199SXin Li
754*9c5db199SXin Li
755*9c5db199SXin Liclass ControlFile(RpcObject):
756*9c5db199SXin Li    """
757*9c5db199SXin Li    AFE control file object
758*9c5db199SXin Li
759*9c5db199SXin Li    Fields: synch_count, dependencies, control_file, is_server
760*9c5db199SXin Li    """
761*9c5db199SXin Li    def __repr__(self):
762*9c5db199SXin Li        return 'CONTROL FILE: %s' % self.control_file
763*9c5db199SXin Li
764*9c5db199SXin Li
765*9c5db199SXin Liclass Label(RpcObject):
766*9c5db199SXin Li    """
767*9c5db199SXin Li    AFE label object
768*9c5db199SXin Li
769*9c5db199SXin Li    Fields:
770*9c5db199SXin Li        name, invalid, platform, kernel_config, id, only_if_needed
771*9c5db199SXin Li    """
772*9c5db199SXin Li    def __repr__(self):
773*9c5db199SXin Li        return 'LABEL: %s' % self.name
774*9c5db199SXin Li
775*9c5db199SXin Li
776*9c5db199SXin Li    def add_hosts(self, hosts):
777*9c5db199SXin Li        # We must use the label's name instead of the id because label ids are
778*9c5db199SXin Li        # not consistent across main-shard.
779*9c5db199SXin Li        return self.afe.run('label_add_hosts', id=self.name, hosts=hosts)
780*9c5db199SXin Li
781*9c5db199SXin Li
782*9c5db199SXin Li    def remove_hosts(self, hosts):
783*9c5db199SXin Li        # We must use the label's name instead of the id because label ids are
784*9c5db199SXin Li        # not consistent across main-shard.
785*9c5db199SXin Li        return self.afe.run('label_remove_hosts', id=self.name, hosts=hosts)
786*9c5db199SXin Li
787*9c5db199SXin Li
788*9c5db199SXin Liclass Acl(RpcObject):
789*9c5db199SXin Li    """
790*9c5db199SXin Li    AFE acl object
791*9c5db199SXin Li
792*9c5db199SXin Li    Fields:
793*9c5db199SXin Li        users, hosts, description, name, id
794*9c5db199SXin Li    """
795*9c5db199SXin Li    def __repr__(self):
796*9c5db199SXin Li        return 'ACL: %s' % self.name
797*9c5db199SXin Li
798*9c5db199SXin Li
799*9c5db199SXin Li    def add_hosts(self, hosts):
800*9c5db199SXin Li        self.afe.log('Adding hosts %s to ACL %s' % (hosts, self.name))
801*9c5db199SXin Li        return self.afe.run('acl_group_add_hosts', self.id, hosts)
802*9c5db199SXin Li
803*9c5db199SXin Li
804*9c5db199SXin Li    def remove_hosts(self, hosts):
805*9c5db199SXin Li        self.afe.log('Removing hosts %s from ACL %s' % (hosts, self.name))
806*9c5db199SXin Li        return self.afe.run('acl_group_remove_hosts', self.id, hosts)
807*9c5db199SXin Li
808*9c5db199SXin Li
809*9c5db199SXin Li    def add_users(self, users):
810*9c5db199SXin Li        self.afe.log('Adding users %s to ACL %s' % (users, self.name))
811*9c5db199SXin Li        return self.afe.run('acl_group_add_users', id=self.name, users=users)
812*9c5db199SXin Li
813*9c5db199SXin Li
814*9c5db199SXin Liclass Job(RpcObject):
815*9c5db199SXin Li    """
816*9c5db199SXin Li    AFE job object
817*9c5db199SXin Li
818*9c5db199SXin Li    Fields:
819*9c5db199SXin Li        name, control_file, control_type, synch_count, reboot_before,
820*9c5db199SXin Li        run_verify, priority, email_list, created_on, dependencies,
821*9c5db199SXin Li        timeout, owner, reboot_after, id
822*9c5db199SXin Li    """
823*9c5db199SXin Li    def __repr__(self):
824*9c5db199SXin Li        return 'JOB: %s' % self.id
825*9c5db199SXin Li
826*9c5db199SXin Li
827*9c5db199SXin Liclass JobStatus(RpcObject):
828*9c5db199SXin Li    """
829*9c5db199SXin Li    AFE job_status object
830*9c5db199SXin Li
831*9c5db199SXin Li    Fields:
832*9c5db199SXin Li        status, complete, deleted, meta_host, host, active, execution_subdir, id
833*9c5db199SXin Li    """
834*9c5db199SXin Li    def __init__(self, afe, hash):
835*9c5db199SXin Li        super(JobStatus, self).__init__(afe, hash)
836*9c5db199SXin Li        self.job = Job(afe, self.job)
837*9c5db199SXin Li        if getattr(self, 'host'):
838*9c5db199SXin Li            self.host = Host(afe, self.host)
839*9c5db199SXin Li
840*9c5db199SXin Li
841*9c5db199SXin Li    def __repr__(self):
842*9c5db199SXin Li        if self.host and self.host.hostname:
843*9c5db199SXin Li            hostname = self.host.hostname
844*9c5db199SXin Li        else:
845*9c5db199SXin Li            hostname = 'None'
846*9c5db199SXin Li        return 'JOB STATUS: %s-%s' % (self.job.id, hostname)
847*9c5db199SXin Li
848*9c5db199SXin Li
849*9c5db199SXin Liclass SpecialTask(RpcObject):
850*9c5db199SXin Li    """
851*9c5db199SXin Li    AFE special task object
852*9c5db199SXin Li    """
853*9c5db199SXin Li    def __init__(self, afe, hash):
854*9c5db199SXin Li        super(SpecialTask, self).__init__(afe, hash)
855*9c5db199SXin Li        self.host = Host(afe, self.host)
856*9c5db199SXin Li
857*9c5db199SXin Li
858*9c5db199SXin Li    def __repr__(self):
859*9c5db199SXin Li        return 'SPECIAL TASK: %s' % self.id
860*9c5db199SXin Li
861*9c5db199SXin Li
862*9c5db199SXin Liclass Host(RpcObject):
863*9c5db199SXin Li    """
864*9c5db199SXin Li    AFE host object
865*9c5db199SXin Li
866*9c5db199SXin Li    Fields:
867*9c5db199SXin Li        status, lock_time, locked_by, locked, hostname, invalid,
868*9c5db199SXin Li        labels, platform, protection, dirty, id
869*9c5db199SXin Li    """
870*9c5db199SXin Li    def __repr__(self):
871*9c5db199SXin Li        return 'HOST OBJECT: %s' % self.hostname
872*9c5db199SXin Li
873*9c5db199SXin Li
874*9c5db199SXin Li    def show(self):
875*9c5db199SXin Li        labels = list(set(self.labels) - set([self.platform]))
876*9c5db199SXin Li        print('%-6s %-7s %-7s %-16s %s' % (self.hostname, self.status,
877*9c5db199SXin Li                                           self.locked, self.platform,
878*9c5db199SXin Li                                           ', '.join(labels)))
879*9c5db199SXin Li
880*9c5db199SXin Li
881*9c5db199SXin Li    def delete(self):
882*9c5db199SXin Li        return self.afe.run('delete_host', id=self.id)
883*9c5db199SXin Li
884*9c5db199SXin Li
885*9c5db199SXin Li    def modify(self, **dargs):
886*9c5db199SXin Li        return self.afe.run('modify_host', id=self.id, **dargs)
887*9c5db199SXin Li
888*9c5db199SXin Li
889*9c5db199SXin Li    def get_acls(self):
890*9c5db199SXin Li        return self.afe.get_acls(hosts__hostname=self.hostname)
891*9c5db199SXin Li
892*9c5db199SXin Li
893*9c5db199SXin Li    def add_acl(self, acl_name):
894*9c5db199SXin Li        self.afe.log('Adding ACL %s to host %s' % (acl_name, self.hostname))
895*9c5db199SXin Li        return self.afe.run('acl_group_add_hosts', id=acl_name,
896*9c5db199SXin Li                            hosts=[self.hostname])
897*9c5db199SXin Li
898*9c5db199SXin Li
899*9c5db199SXin Li    def remove_acl(self, acl_name):
900*9c5db199SXin Li        self.afe.log('Removing ACL %s from host %s' % (acl_name, self.hostname))
901*9c5db199SXin Li        return self.afe.run('acl_group_remove_hosts', id=acl_name,
902*9c5db199SXin Li                            hosts=[self.hostname])
903*9c5db199SXin Li
904*9c5db199SXin Li
905*9c5db199SXin Li    def get_labels(self):
906*9c5db199SXin Li        return self.afe.get_labels(host__hostname__in=[self.hostname])
907*9c5db199SXin Li
908*9c5db199SXin Li
909*9c5db199SXin Li    def add_labels(self, labels):
910*9c5db199SXin Li        self.afe.log('Adding labels %s to host %s' % (labels, self.hostname))
911*9c5db199SXin Li        return self.afe.run('host_add_labels', id=self.id, labels=labels)
912*9c5db199SXin Li
913*9c5db199SXin Li
914*9c5db199SXin Li    def remove_labels(self, labels):
915*9c5db199SXin Li        self.afe.log('Removing labels %s from host %s' % (labels,self.hostname))
916*9c5db199SXin Li        return self.afe.run('host_remove_labels', id=self.id, labels=labels)
917*9c5db199SXin Li
918*9c5db199SXin Li
919*9c5db199SXin Li    def is_available(self):
920*9c5db199SXin Li        """Check whether DUT host is available.
921*9c5db199SXin Li
922*9c5db199SXin Li        @return: bool
923*9c5db199SXin Li        """
924*9c5db199SXin Li        return not (self.locked
925*9c5db199SXin Li                    or self.status in host_states.UNAVAILABLE_STATES)
926*9c5db199SXin Li
927*9c5db199SXin Li
928*9c5db199SXin Liclass User(RpcObject):
929*9c5db199SXin Li    def __repr__(self):
930*9c5db199SXin Li        return 'USER: %s' % self.login
931*9c5db199SXin Li
932*9c5db199SXin Li
933*9c5db199SXin Liclass TestStatus(RpcObject):
934*9c5db199SXin Li    """
935*9c5db199SXin Li    TKO test status object
936*9c5db199SXin Li
937*9c5db199SXin Li    Fields:
938*9c5db199SXin Li        test_idx, hostname, testname, id
939*9c5db199SXin Li        complete_count, incomplete_count, group_count, pass_count
940*9c5db199SXin Li    """
941*9c5db199SXin Li    def __repr__(self):
942*9c5db199SXin Li        return 'TEST STATUS: %s' % self.id
943*9c5db199SXin Li
944*9c5db199SXin Li
945*9c5db199SXin Liclass HostAttribute(RpcObject):
946*9c5db199SXin Li    """
947*9c5db199SXin Li    AFE host attribute object
948*9c5db199SXin Li
949*9c5db199SXin Li    Fields:
950*9c5db199SXin Li        id, host, attribute, value
951*9c5db199SXin Li    """
952*9c5db199SXin Li    def __repr__(self):
953*9c5db199SXin Li        return 'HOST ATTRIBUTE %d' % self.id
954