Add: HttpAPI class for using REST-ful HTTP API

4 files changed, 270 insertions, 17 deletions

diff --git a/ick2/__init__.py b/ick2/__init__.py
index f5770f6..257fa2a 100644
--- a/ick2/__init__.py
+++ b/ick2/__init__.py
@@ -51,3 +51,5 @@ from .controllerapi import (
 )
 from .blobapi import BlobAPI
 from .blob_store import BlobStore
+
+from .client import HttpAPI, HttpError
diff --git a/ick2/client.py b/ick2/client.py
new file mode 100644
index 0000000..a474987
--- /dev/null
+++ b/ick2/client.py
@@ -0,0 +1,105 @@
+# Copyright (C) 2018  Lars Wirzenius
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+
+import json
+
+
+import requests
+
+
+class HttpError(Exception):
+
+    pass
+
+
+class HttpAPI:
+
+    # Make requests to an HTTP API.
+
+    json_type = 'application/json'
+
+    def __init__(self):
+        self._session = requests.Session()
+        self._token = None
+
+    def set_session(self, session):
+        self._session = session
+
+    def set_token(self, token):
+        self._token = token
+
+    def get_dict(self, url, headers=None):
+        r = self._request(self._session.get, url, headers=headers)
+        ct = r.headers.get('Content-Type')
+        if ct != self.json_type:
+            raise HttpError('Not JSON response')
+        try:
+            return r.json()
+        except json.decoder.JSONDecodeError:
+            raise HttpError('JSON parsing error')
+
+    def get_blob(self, url, headers=None):
+        r = self._request(self._session.get, url, headers=headers)
+        return r.content
+
+    def post(self, url, headers=None, body=None):
+        self._send_request(self._session.post, url, headers=headers, body=body)
+        return None
+
+    def put(self, url, headers=None, body=None):
+        self._send_request(self._session.put, url, headers=headers, body=body)
+        return None
+
+    def _send_request(self, func, url, headers=None, body=None):
+        if headers is None:
+            headers = {}
+        headers = dict(headers)
+        h, body = self._get_content_type_header(body)
+        headers.update(h)
+        self._request(func, url, headers=headers, data=body)
+        return None
+
+    def _get_content_type_header(self, body):
+        if isinstance(body, dict):
+            header = {
+                'Content-Type': 'application/json',
+            }
+            body = json.dumps(body)
+            return header, body
+        return {}, body
+
+    def _get_authorization_headers(self):
+        return {
+            'Authorization': 'Bearer {}'.format(self._token),
+        }
+
+    def _request(self, func, url, headers=None, **kwargs):
+        if headers is None:
+            headers = {}
+        headers.update(self._get_authorization_headers())
+
+        logging.debug('request: func=%r', func)
+        logging.debug('request: url=%r', url)
+        for h in headers:
+            logging.debug('request: %s: %s', h, headers[h])
+        logging.debug('request: kwargs=%r', kwargs)
+
+        r = func(url, headers=headers, verify=False, **kwargs)
+        logging.debug('response: status_code=%r', r.status_code)
+        logging.debug('response: content=%r', r.content)
+
+        if not r.ok:
+            raise HttpError(r.status_code)
+        return r
diff --git a/ick2/client_tests.py b/ick2/client_tests.py
new file mode 100644
index 0000000..1a0bcc5
--- /dev/null
+++ b/ick2/client_tests.py
@@ -0,0 +1,146 @@
+# Copyright (C) 2018  Lars Wirzenius
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+
+import json
+import unittest
+
+
+import ick2
+
+
+json_type = 'application/json'
+
+
+class HttpAPITests(unittest.TestCase):
+
+    def setUp(self):
+        self.session = FakeHttpSession()
+        self.session.token = 'SECRET-TOKEN'
+        self.client = ick2.HttpAPI()
+        self.client.set_session(self.session)
+        self.client.set_token(self.session.token)
+
+    def test_get_dict_raises_exception_on_error(self):
+        self.session.response = FakeResponse(400)
+        with self.assertRaises(ick2.HttpError):
+            self.client.get_dict('http://controller/version')
+
+    def test_get_dict_raises_exception_on_not_json(self):
+        self.session.response = FakeResponse(200, body=json.dumps('{}'))
+        with self.assertRaises(ick2.HttpError):
+            self.client.get_dict('http://controller/version')
+
+    def test_get_dict_raises_exception_on_malformed_json(self):
+        self.session.response = FakeResponse(
+            200, body='this is not really JSON', content_type=json_type)
+        with self.assertRaises(ick2.HttpError):
+            self.client.get_dict('http://controller/version')
+
+    def test_get_dict_returns_response(self):
+        version = {
+            'version': '1.0',
+        }
+
+        self.session.response = FakeResponse(
+            200, body=json.dumps(version), content_type=json_type)
+        obj = self.client.get_dict('http://controller/version')
+        self.assertEqual(obj, version)
+
+    def test_get_blob_raises_exception_on_error(self):
+        self.session.response = FakeResponse(404)
+        with self.assertRaises(ick2.HttpError):
+            self.client.get_blob('http://blobs/blob/foo')
+
+    def test_get_blob_returns_response(self):
+        blob = b'hello, world\n'
+        self.session.response = FakeResponse(200, body=blob)
+        obj = self.client.get_blob('http://blobs/blob/foo')
+        self.assertEqual(obj, blob)
+
+    def test_post_raises_exception_on_error(self):
+        self.session.response = FakeResponse(400)
+        with self.assertRaises(ick2.HttpError):
+            self.client.post('http://controller/work', body='')
+
+    def test_post_succeeds(self):
+        work = {
+            'project': 'foo',
+            'stdout': 'hello, world\n',
+        }
+
+        self.session.response = FakeResponse(
+            201, body=json.dumps(work), content_type=json_type)
+        obj = self.client.post('http://controller/work', body=work)
+        self.assertEqual(obj, None)
+
+    def test_put_raises_exception_on_error(self):
+        blob = b'fooblob'
+        self.session.response = FakeResponse(400)
+        with self.assertRaises(ick2.HttpError):
+            self.client.put('http://blobs/blob/foo', body=blob)
+
+    def test_put_succeeds(self):
+        blob = b'fooblob'
+        self.session.response = FakeResponse(201, body=None)
+        obj = self.client.put('http://controller/work', body=blob)
+        self.assertEqual(obj, None)
+
+
+class FakeHttpSession:
+
+    def __init__(self):
+        self.response = None
+        self.token = None
+
+    def get(self, url, headers=None, verify=None):
+        assert self.response is not None
+        assert self.is_authorized(headers)
+        return self.response
+
+    def post(self, url, headers=None, data=None, verify=None):
+        assert self.response is not None
+        assert self.is_authorized(headers)
+        return self.response
+
+    def put(self, url, headers=None, data=None, verify=None):
+        assert self.response is not None
+        assert self.is_authorized(headers)
+        return self.response
+
+    def is_authorized(self, headers):
+        assert self.token is not None
+        assert 'Authorization' in headers
+        v = headers['Authorization']
+        return v == 'Bearer {}'.format(self.token)
+
+
+class FakeResponse:
+
+    def __init__(self, status_code, body=None, content_type=None):
+        if content_type is None:
+            content_type = 'application/octet-stream'
+
+        self.status_code = status_code
+        self.headers = {
+            'Content-Type': content_type,
+        }
+        self.content = body
+
+    @property
+    def ok(self):
+        return self.status_code in (200, 201)
+
+    def json(self):
+        return json.loads(self.content)
diff --git a/worker_manager b/worker_manager
index 7530469..db4d08f 100755
--- a/worker_manager
+++ b/worker_manager
@@ -26,7 +26,6 @@ import time
 import apifw
 import cliapp
 import Crypto.PublicKey.RSA
-import requests
 import urllib3
 import ick2
 
@@ -193,31 +192,32 @@ class ControllerAPI:
 class HttpApi:
 
     def __init__(self):
-        self._session = requests.Session()
+        self._httpapi = ick2.HttpAPI()
 
     def post(self, url, headers, body):
-        r = self._session.post(url, json=body, headers=headers, verify=False)
-        if not r.ok:
-            logging.warning('Error: POST %s returned %s', url, r.status_code)
-        return r.status_code
+        ret = self._request(self._httpapi.post, url, headers, body=body)
+        if ret is None:
+            return 400
+        return 201
 
     def put(self, url, headers, body):
-        r = self._session.put(url, data=body, headers=headers, verify=False)
-        if not r.ok:
-            logging.warning('Error: PUT %s returned %s', url, r.status_code)
-        return r.status_code
+        ret = self._request(self._httpapi.put, url, headers, body=body)
+        if ret is None:
+            return 400
+        return 200
 
     def get(self, url, headers):
-        r = self._session.get(url, headers=headers, verify=False)
-        if not r.ok or not r.text:
-            return None
-        return r.json()
+        return self._request(self._httpapi.get_dict, url, headers)
 
     def get_blob(self, url, headers):
-        r = self._session.get(url, headers=headers, verify=False)
-        if not r.ok:
+        return self._request(self._httpapi.get_blob, url, headers)
+
+    def _request(self, func, url, headers, **kwargs):
+        try:
+            return func(url, headers=headers, **kwargs)
+        except ick2.HttpError as e:
+            logging.warning('Error: %s returned %s', url, e)
             return None
-        return r.content
 
 
 class TokenGenerator: