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.server import utils
14from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
15
16AFE = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
17
18
19def host_in_lab(host):
20    """Check if the host is in the lab and an object the AFE knows.
21
22    This check ensures that autoserv and the host's current job is running
23    inside a fully Autotest instance, aka a lab environment. If this is the
24    case it then verifies the host is registed with the configured AFE
25    instance.
26
27    @param host: Host object to verify.
28
29    @returns The host model object.
30    """
31    if not host.job or not host.job.in_lab:
32        return False
33    return AFE.get_hosts(hostname=host.hostname)
34
35
36def get_build(host):
37    """Retrieve the current build for a given host from the AFE.
38
39    Looks through a host's labels in the AFE to determine its build.
40
41    @param host: Host object to get build.
42
43    @returns The current build or None if it could not find it or if there
44             were multiple build labels assigned to the host.
45    """
46    if not host_in_lab(host):
47        return None
48    return utils.get_build_from_afe(host.hostname, AFE)
49
50
51def get_board(host):
52    """Retrieve the board for a given host from the AFE.
53
54    Contacts the AFE to retrieve the board for a given host.
55
56    @param host: Host object to get board.
57
58    @returns The current board or None if it could not find it.
59    """
60    if not host_in_lab(host):
61        return None
62    return utils.get_board_from_afe(host.hostname, AFE)
63
64
65def clear_version_labels(host):
66    """Clear version labels for a given host.
67
68    @param host: Host whose version labels to clear.
69    """
70    if not host_in_lab(host):
71        return
72
73    host_list = [host.hostname]
74    labels = AFE.get_labels(
75            name__startswith=host.VERSION_PREFIX,
76            host__hostname=host.hostname)
77
78    for label in labels:
79        label.remove_hosts(hosts=host_list)
80
81
82def add_version_label(host, image_name):
83    """Add version labels to a host.
84
85    @param host: Host to add the version label for.
86    @param image_name: Name of the build version to add to the host.
87    """
88    if not host_in_lab(host):
89        return
90    label = '%s:%s' % (host.VERSION_PREFIX, image_name)
91    AFE.run('label_add_hosts', id=label, hosts=[host.hostname])
92
93
94def machine_install_and_update_labels(host, *args, **dargs):
95    """Calls machine_install and updates the version labels on a host.
96
97    @param host: Host object to run machine_install on.
98    @param *args: Args list to pass to machine_install.
99    @param **dargs: dargs dict to pass to machine_install.
100    """
101    clear_version_labels(host)
102    image_name = host.machine_install(*args, **dargs)
103    add_version_label(host, image_name)
104
105
106def get_stable_version(board, android=False):
107    """Retrieves a board's stable version from the AFE.
108
109    @param board: Board to lookup.
110    @param android: If True, indicates we are looking up a Android/Brillo-based
111                    board. There is no default version that works for all
112                    Android/Brillo boards. If False, we are looking up a Chrome
113                    OS based board.
114
115    @returns Stable version of the given board.
116    """
117    return AFE.run('get_stable_version', board=board, android=android)
118