afe_utils.py revision b3b6db3d6f9e7805fd8209317c51f4fcc3b835ac 
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
7import common
8from autotest_lib.server import utils
9from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
10
11AFE = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
12
13
14def host_in_lab(host):
15    """Check if the host is in the lab and an object the AFE knows.
16
17    This check ensures that autoserv and the host's current job is running
18    inside a fully Autotest instance, aka a lab environment. If this is the
19    case it then verifies the host is registed with the configured AFE
20    instance.
21
22    @param host: Host object to verify.
23
24    @returns The host model object.
25    """
26    if not host.job.in_lab:
27        return False
28    return AFE.get_hosts(hostname=host.hostname)
29
30
31def get_build(host):
32    """Retrieve the current build for a given host from the AFE.
33
34    Looks through a host's labels in the AFE to determine its build.
35
36    @param host: Host object to get build.
37
38    @returns The current build or None if it could not find it or if there
39             were multiple build labels assigned to the host.
40    """
41    if not host_in_lab(host):
42        return None
43    return utils.get_build_from_afe(host.hostname, AFE)
44
45
46def clear_version_labels(host):
47    """Clear version labels for a given host.
48
49    @param host: Host whose version labels to clear.
50    """
51    if not host_in_lab(host):
52        return
53
54    host_list = [host.hostname]
55    labels = AFE.get_labels(
56            name__startswith=host.VERSION_PREFIX,
57            host__hostname=host.hostname)
58
59    for label in labels:
60        label.remove_hosts(hosts=host_list)
61
62
63def add_version_label(host, image_name):
64    """Add version labels to a host.
65
66    @param host: Host to add the version label for.
67    @param image_name: Name of the build version to add to the host.
68    """
69    if not host_in_lab(host):
70        return
71    label = '%s:%s' % (host.VERSION_PREFIX, image_name)
72    AFE.run('label_add_hosts', id=label, hosts=[host.hostname])
73
74
75def machine_install_and_update_labels(host, *args, **dargs):
76    """Calls machine_install and updates the version labels on a host.
77
78    @param host: Host object to run machine_install on.
79    @param *args: Args list to pass to machine_install.
80    @param **dargs: dargs dict to pass to machine_install.
81    """
82    clear_version_labels(host)
83    image_name = host.machine_install(*args, **dargs)
84    add_version_label(host, image_name)
85