Add support for OAuth2 Bearer tokens

tes/client.py

@@ -29,6 +29,9 @@ class HTTPClient(object):
     password = attrib(default=None,
                       convert=strconv,
                       validator=optional(instance_of(str)))
+    token = attrib(default=None,
+                   convert=strconv,
+                   validator=optional(instance_of(str)))
 
     @url.validator
     def __check_url(self, attribute, value):
@@ -40,11 +43,10 @@ def __check_url(self, attribute, value):
             )
 
     def get_service_info(self):
+        kwargs = self._request_params()
         response = requests.get(
             "%s/v1/tasks/service-info" % (self.url),
-            timeout=self.timeout,
-            auth=(self.user, self.password)
-        )
+            **kwargs)
         raise_for_status(response)
         return unmarshal(response.json(), ServiceInfo)
 
@@ -53,35 +55,31 @@ def create_task(self, task):
             msg = task.as_json()
         else:
             raise TypeError("Expected Task instance")
+
+        kwargs = self._request_params(data=msg)
         response = requests.post(
             "%s/v1/tasks" % (self.url),
-            data=msg,
-            headers={'Content-Type': 'application/json'},
-            timeout=self.timeout,
-            auth=(self.user, self.password)
+            **kwargs
         )
         raise_for_status(response)
         return unmarshal(response.json(), CreateTaskResponse).id
 
     def get_task(self, task_id, view="BASIC"):
         req = GetTaskRequest(task_id, view)
         payload = {"view": req.view}
+        kwargs = self._request_params(params=payload)
         response = requests.get(
             "%s/v1/tasks/%s" % (self.url, req.id),
-            params=payload,
-            timeout=self.timeout,
-            auth=(self.user, self.password)
-        )
+            **kwargs)
         raise_for_status(response)
         return unmarshal(response.json(), Task)
 
     def cancel_task(self, task_id):
         req = CancelTaskRequest(task_id)
+        kwargs = self._request_params()
         response = requests.post(
             "%s/v1/tasks/%s:cancel" % (self.url, req.id),
-            timeout=self.timeout,
-            auth=(self.user, self.password)
-        )
+            **kwargs)
         raise_for_status(response)
         return
 
@@ -95,12 +93,10 @@ def list_tasks(self, view="MINIMAL", page_size=None, page_token=None):
         )
         msg = req.as_dict()
 
+        kwargs = self._request_params(params=msg)
         response = requests.get(
             "%s/v1/tasks" % (self.url),
-            params=msg,
-            timeout=self.timeout,
-            auth=(self.user, self.password)
-        )
+            **kwargs)
         raise_for_status(response)
         return unmarshal(response.json(), ListTasksResponse)
 
@@ -122,5 +118,20 @@ def check_success(data):
 
             if max_time is not None and time.time() >= max_time:
                 raise TimeoutError("last_response: %s" % (response.as_dict()))
-
             time.sleep(0.5)
+
+    def _request_params(self, data=None, params=None):
+        kwargs = {'timeout': self.timeout}
+
+        if data:
+            kwargs['data'] = data
+        if params:
+            kwargs['params'] = params
+        if self.token:
+            kwargs['headers'] = {'Content-type': 'application/json',
+                                 'Authorization': 'Bearer ' + self.token}
+        else:
+            kwargs['headers'] = {'Content-type': 'application/json'}
+            kwargs['auth'] = (self.user, self.password)
+
+        return kwargs