Add: ControllerClient class for using controller API

1 files changed, 9 insertions, 22 deletions

diff --git a/worker_manager b/worker_manager
index db4d08f..ec8c19f 100755
--- a/worker_manager
+++ b/worker_manager
@@ -122,34 +122,21 @@ class ControllerAPI:
         self._blob_url = None
         self._token_generator = token_generator
         self._httpapi = HttpApi()
+        self._cc = ick2.ControllerClient()
+        self._cc.set_client_name(name)
+        self._cc.set_controller_url(url)
 
     def register(self):
-        logging.info('Registering worker %s to %s', self._name, self._url)
-        url = self.url('/workers')
-        headers = self.get_auth_headers()
-        body = {
-            'worker': self._name,
-        }
-        code = self._httpapi.post(url, headers, body)
-        if code not in [200, 201, 409]:
-            raise cliapp.AppException('Failed to register worker')
+        self._cc.set_token(self._token_generator.get_token())
+        self._cc.register()
 
     def get_work(self):
-        url = self.url('/work/{}'.format(self._name))
-        headers = self.get_auth_headers()
-        work = self._httpapi.get(url, headers)
-        if work:
-            logging.info('Response: %r', work)
-        return work
+        self._cc.set_token(self._token_generator.get_token())
+        return self._cc.get_work()
 
     def report_work(self, work):
-        logging.info('POST %s', work)
-        url = self.url('/work')
-        headers = self.get_auth_headers()
-        code = self._httpapi.post(url, headers, work)
-        if code not in (200, 201):
-            raise cliapp.AppException(
-                'Error posting data to controller: {}'.format(code))
+        self._cc.set_token(self._token_generator.get_token())
+        self._cc.report_work(work)
 
     def url(self, path):
         return '{}{}'.format(self._url, path)