• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2016 The Chromium 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"""Utility functions for AFE-based interactions.
6
7NOTE: This module should only be used in the context of a running test. Any
8      utilities that require accessing the AFE, should do so by creating
9      their own instance of the AFE client and interact with it directly.
10"""
11
12import common
13from autotest_lib.client.common_lib import error
14from autotest_lib.client.common_lib import global_config
15from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
16
17
18AFE = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
19_CROS_VERSION_MAP = AFE.get_stable_version_map(AFE.CROS_IMAGE_TYPE)
20_FIRMWARE_VERSION_MAP = AFE.get_stable_version_map(AFE.FIRMWARE_IMAGE_TYPE)
21_FAFT_VERSION_MAP = AFE.get_stable_version_map(AFE.FAFT_IMAGE_TYPE)
22_ANDROID_VERSION_MAP = AFE.get_stable_version_map(AFE.ANDROID_IMAGE_TYPE)
23
24_CONFIG = global_config.global_config
25ENABLE_DEVSERVER_TRIGGER_AUTO_UPDATE = _CONFIG.get_config_value(
26        'CROS', 'enable_devserver_trigger_auto_update', type=bool,
27        default=False)
28
29
30def _host_in_lab(host):
31    """Check if the host is in the lab and an object the AFE knows.
32
33    This check ensures that autoserv and the host's current job is running
34    inside a fully Autotest instance, aka a lab environment. If this is the
35    case it then verifies the host is registed with the configured AFE
36    instance.
37
38    @param host: Host object to verify.
39
40    @returns The host model object.
41    """
42    if not host.job or not host.job.in_lab:
43        return False
44    return host._afe_host
45
46
47def get_labels(host, prefix=None):
48    """Get labels of a host with name started with given prefix.
49
50    @param prefix: Prefix of label names, if None, return all labels.
51
52    @returns List of labels that match the prefix or if prefix is None, all
53             labels.
54    """
55    if not prefix:
56        return host._afe_host.labels
57
58    return [label for label in host._afe_host.labels
59            if label.startswith(prefix)]
60
61
62def clear_version_labels(host):
63    """Clear version labels for a given host.
64
65    @param host: Host whose version labels to clear.
66    """
67    host._afe_host.labels = [label for label in host._afe_host.labels
68                             if not label.startswith(host.VERSION_PREFIX)]
69    if not _host_in_lab(host):
70        return
71
72    host_list = [host.hostname]
73    labels = AFE.get_labels(
74            name__startswith=host.VERSION_PREFIX,
75            host__hostname=host.hostname)
76
77    for label in labels:
78        label.remove_hosts(hosts=host_list)
79
80
81def add_version_label(host, image_name):
82    """Add version labels to a host.
83
84    @param host: Host to add the version label for.
85    @param image_name: Name of the build version to add to the host.
86    """
87    label = '%s:%s' % (host.VERSION_PREFIX, image_name)
88    host._afe_host.labels.append(label)
89    if not _host_in_lab(host):
90        return
91    AFE.run('label_add_hosts', id=label, hosts=[host.hostname])
92
93
94def get_stable_cros_image_name(board):
95    """Retrieve the Chrome OS stable image name for a given board.
96
97    @param board: Board to lookup.
98
99    @returns Name of a Chrome OS image to be installed in order to
100            repair the given board.
101    """
102    return _CROS_VERSION_MAP.get_image_name(board)
103
104
105def get_stable_firmware_version(board):
106    """Retrieve the stable firmware version for a given board.
107
108    @param board: Board to lookup.
109
110    @returns A version of firmware to be installed via
111             `chromeos-firmwareupdate` from a repair build.
112    """
113    return _FIRMWARE_VERSION_MAP.get_version(board)
114
115
116def get_stable_faft_version(board):
117    """Retrieve the stable firmware version for FAFT DUTs.
118
119    @param board: Board to lookup.
120
121    @returns A version of firmware to be installed in order to
122            repair firmware on a DUT used for FAFT testing.
123    """
124    return _FAFT_VERSION_MAP.get_version(board)
125
126
127def get_stable_android_version(board):
128    """Retrieve the stable Android version a given board.
129
130    @param board: Board to lookup.
131
132    @returns Stable version of Android for the given board.
133    """
134    return _ANDROID_VERSION_MAP.get_version(board)
135
136
137def get_host_attribute(host, attribute, use_local_value=True):
138    """Looks up the value of host attribute for the host.
139
140    @param host: A Host object to lookup for attribute value.
141    @param attribute: Name of the host attribute.
142    @param use_local_value: Boolean to indicate if the local value or AFE value
143            should be retrieved.
144
145    @returns value for the given attribute or None if not found.
146    """
147    local_value = host._afe_host.attributes.get(attribute)
148    if not _host_in_lab(host) or use_local_value:
149        return local_value
150
151    hosts = AFE.get_hosts(hostname=host.hostname)
152    if hosts and attribute in hosts[0].attributes:
153        return hosts[0].attributes[attribute]
154    else:
155        return local_value
156
157
158def _clear_host_attributes_before_provision(host, info):
159    """Clear host attributes before provision, e.g., job_repo_url.
160
161    @param host: A Host object to clear attributes before provision.
162    @param info: A HostInfo to update the attributes in.
163    """
164    attributes = host.get_attributes_to_clear_before_provision()
165    if not attributes:
166        return
167
168    for key in attributes:
169        info.attributes.pop(key, None)
170
171
172def update_host_attribute(host, attribute, value):
173    """Updates the host attribute with given value.
174
175    @param host: A Host object to update attribute value.
176    @param attribute: Name of the host attribute.
177    @param value: Value for the host attribute.
178
179    @raises AutoservError: If we failed to update the attribute.
180    """
181    host._afe_host.attributes[attribute] = value
182    if not _host_in_lab(host):
183        return
184
185    AFE.set_host_attribute(attribute, value, hostname=host.hostname)
186    info = host.host_info_store.get(force_refresh=True)
187    if info.attributes.get(attribute) != value:
188        raise error.AutoservError(
189                'Failed to update host attribute `%s` with %s, host %s' %
190                (attribute, value, host.hostname))
191
192
193def machine_install_and_update_labels(host, *args, **dargs):
194    """Calls machine_install and updates the version labels on a host.
195
196    @param host: Host object to run machine_install on.
197    @param *args: Args list to pass to machine_install.
198    @param **dargs: dargs dict to pass to machine_install.
199    """
200    info = host.host_info_store.get()
201    info.clear_version_labels()
202    _clear_host_attributes_before_provision(host, info)
203    host.host_info_store.commit(info)
204    # If ENABLE_DEVSERVER_TRIGGER_AUTO_UPDATE is enabled and the host is a
205    # CrosHost, devserver will be used to trigger auto-update.
206    if host.support_devserver_provision:
207        image_name, host_attributes = host.machine_install_by_devserver(
208            *args, **dargs)
209    else:
210        image_name, host_attributes = host.machine_install(*args, **dargs)
211
212    info = host.host_info_store.get()
213    info.attributes.update(host_attributes)
214    info.set_version_label(host.VERSION_PREFIX, image_name)
215    host.host_info_store.commit(info)
216