From a692650307e06a5c93c46d6c9f28f93cd473293e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juha=20T=C3=B6rnroos?= <juha.tornroos@csc.fi>
Date: Mon, 24 Sep 2018 18:48:01 +0300
Subject: [PATCH 1/2] Add support for OAuth2 Bearer tokens

---
 tes/client.py | 65 ++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 46 insertions(+), 19 deletions(-)

diff --git a/tes/client.py b/tes/client.py
index 980a9d9..8263ced 100644
--- a/tes/client.py
+++ b/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,47 @@ 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 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
+
     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 +109,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 +134,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

From 1cac35c9b8d854b076d9272b252357a73f447e19 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juha=20T=C3=B6rnroos?= <juha.tornroos@csc.fi>
Date: Mon, 24 Sep 2018 19:23:05 +0300
Subject: [PATCH 2/2] Remove a duplicate method

---
 tes/client.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/tes/client.py b/tes/client.py
index 8263ced..c9e0345 100644
--- a/tes/client.py
+++ b/tes/client.py
@@ -64,22 +64,6 @@ def create_task(self, task):
         raise_for_status(response)
         return unmarshal(response.json(), CreateTaskResponse).id
 
-    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
-
     def get_task(self, task_id, view="BASIC"):
         req = GetTaskRequest(task_id, view)
         payload = {"view": req.view}