1#
2# Copyright (C) 2017 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import httplib2
18import logging
19import threading
20import time
21
22from googleapiclient import discovery
23from googleapiclient import http
24from oauth2client.service_account import ServiceAccountCredentials
25
26from host_controller.tfc import command_task
27
28API_NAME = "tradefed_cluster"
29API_VERSION = "v1"
30SCOPES = ['https://www.googleapis.com/auth/userinfo.email']
31
32
33class TfcClient(object):
34    """The class for accessing TFC API.
35
36    Attributes:
37        _service: The TFC service.
38    """
39
40    def __init__(self, service):
41        self._service = service
42
43    def LeaseHostTasks(self, cluster_id, next_cluster_ids, hostname, device_infos):
44        """Calls leasehosttasks.
45
46        Args:
47            cluster_id: A string, the primary cluster to lease tasks from.
48            next_cluster_ids: A list of Strings, the secondary clusters to lease
49                              tasks from.
50            hostname: A string, the name of the TradeFed host.
51            device_infos: A list of DeviceInfo, the information about the
52                          devices connected to the host.
53
54        Returns:
55            A list of command_task.CommandTask, the leased tasks.
56        """
57        lease = {"hostname": hostname,
58                 "cluster": cluster_id,
59                 "next_cluster_ids": next_cluster_ids,
60                 "device_infos": [x.ToLeaseHostTasksJson()
61                                  for x in device_infos]}
62        logging.info("tasks.leasehosttasks body=%s", lease)
63        tasks = self._service.tasks().leasehosttasks(body=lease).execute()
64        logging.info("tasks.leasehosttasks response=%s", tasks)
65        if "tasks" not in tasks:
66            return []
67        return [command_task.CommandTask(**task) for task in tasks["tasks"]]
68
69    def TestResourceList(self, request_id):
70        """Calls testResource.list.
71
72        Args:
73            request_id: int, id of request to grab resources for
74
75        Returns:
76            A list of TestResources
77        """
78        logging.info("request.testResource.list request_id=%s", request_id)
79        test_resources = self._service.requests().testResource().list(request_id=request_id).execute()
80        logging.info("request.testResource.list response=%s", test_resources)
81        if 'test_resources' not in test_resources:
82            return {}
83        return test_resources['test_resources']
84
85    @staticmethod
86    def CreateDeviceSnapshot(cluster_id, hostname, dev_infos):
87        """Creates a DeviceSnapshot which can be uploaded as host event.
88
89        Args:
90            cluster_id: A string, the cluster to upload snapshot to.
91            hostname: A string, the name of the TradeFed host.
92            dev_infos: A list of DeviceInfo.
93
94        Returns:
95            A JSON object.
96        """
97        obj = {"time": int(time.time()),
98               "data": {},
99               "cluster": cluster_id,
100               "hostname": hostname,
101               "tf_version": "(unknown)",
102               "type": "DeviceSnapshot",
103               "device_infos": [x.ToDeviceSnapshotJson() for x in dev_infos]}
104        return obj
105
106    def SubmitHostEvents(self, host_events):
107        """Calls host_events.submit.
108
109        Args:
110            host_events: A list of JSON objects. Currently DeviceSnapshot is
111                         the only type of host events.
112        """
113        json_obj = {"host_events": host_events}
114        logging.info("host_events.submit body=%s", json_obj)
115        self._service.host_events().submit(body=json_obj).execute()
116
117    def SubmitCommandEvents(self, command_events):
118        """Calls command_events.submit.
119
120        Args:
121            command_events: A list of JSON objects converted from CommandAttempt.
122        """
123        json_obj = {"command_events": command_events}
124        logging.info("command_events.submit body=%s", json_obj)
125        self._service.command_events().submit(body=json_obj).execute()
126
127    def NewRequest(self, request):
128        """Calls requests.new.
129
130        Args:
131            request: An instance of Request.
132
133        Returns:
134            A JSON object, the new request queued in the cluster.
135
136            Sample
137            {'state': 'UNKNOWN',
138             'command_line': 'vts-codelab --run-target sailfish',
139             'id': '2',
140             'user': 'testuser'}
141        """
142        body = request.GetBody()
143        params = request.GetParameters()
144        logging.info("requests.new parameters=%s body=%s", params, body)
145        return self._service.requests().new(body=body, **params).execute()
146
147
148def CreateTfcClient(api_root, oauth2_service_json,
149                    api_name=API_NAME, api_version=API_VERSION, scopes=SCOPES):
150    """Builds an object of TFC service from a given URL.
151
152    Args:
153        api_root: The URL to the service.
154        oauth2_service_json: The path to service account key file.
155
156    Returns:
157        A TfcClient object.
158    """
159    discovery_url = "%s/discovery/v1/apis/%s/%s/rest" % (
160            api_root, api_name, api_version)
161    logging.info("Build service from: %s", discovery_url)
162    credentials = ServiceAccountCredentials.from_json_keyfile_name(
163            oauth2_service_json, scopes=scopes)
164    # httplib2.Http is not thread-safe. Use thread local object.
165    thread_local = threading.local()
166    thread_local.http = credentials.authorize(httplib2.Http())
167    def BuildHttpRequest(unused_http, *args, **kwargs):
168        if not hasattr(thread_local, "http"):
169            thread_local.http = credentials.authorize(httplib2.Http())
170        return http.HttpRequest(thread_local.http, *args, **kwargs)
171
172    service = discovery.build(
173            api_name, api_version, http=thread_local.http,
174            discoveryServiceUrl=discovery_url,
175            requestBuilder=BuildHttpRequest)
176    return TfcClient(service)
177