Add: ControllerClient class for using controller API

1 files changed, 56 insertions, 0 deletions

diff --git a/ick2/client.py b/ick2/client.py
index a474987..28d4c1e 100644
--- a/ick2/client.py
+++ b/ick2/client.py
@@ -14,6 +14,7 @@
 
 
 import json
+import logging
 
 
 import requests
@@ -103,3 +104,58 @@ class HttpAPI:
         if not r.ok:
             raise HttpError(r.status_code)
         return r
+
+
+class ControllerClient:
+
+    def __init__(self):
+        self._name = None
+        self._api = HttpAPI()
+        self._url = None
+
+    def set_client_name(self, name):
+        self._name = name
+
+    def set_http_api(self, api):
+        self._api = api
+
+    def set_controller_url(self, url):
+        self._url = url
+
+    def set_token(self, token):
+        self._api.set_token(token)
+
+    def url(self, path):
+        return '{}{}'.format(self._url, path)
+
+    def get_blob_service_url(self):
+        url = self.url('/version')
+        version = self._api.get_dict(url)
+        return version.get('blob_service')
+
+    def register(self):
+        assert self._url is not None
+        url = self.url('/workers')
+        logging.info('Registering worker %s to %s', self._name, url)
+        body = {
+            'worker': self._name,
+        }
+        try:
+            self._api.post(url, body=body)
+        except HttpError:
+            pass
+
+    def get_work(self):
+        url = self.url('/work/{}'.format(self._name))
+        work = self._api.get_dict(url)
+        logging.info('Requested work, got: %r', work)
+        return work
+
+    def report_work(self, work):
+        logging.info('Reporting work: %r', work)
+        url = self.url('/work')
+        headers = {
+            'Content-Type': self._api.json_type,
+        }
+        body = json.dumps(work)
+        self._api.post(url, headers=headers, body=body)