diff options
| -rw-r--r-- | ick2/__init__.py | 2 | ||||
| -rw-r--r-- | ick2/client.py | 56 | ||||
| -rw-r--r-- | ick2/client_tests.py | 67 | ||||
| -rwxr-xr-x | worker_manager | 31 |
4 files changed, 133 insertions, 23 deletions
diff --git a/ick2/__init__.py b/ick2/__init__.py index 257fa2a..0fd0013 100644 --- a/ick2/__init__.py +++ b/ick2/__init__.py @@ -52,4 +52,4 @@ from .controllerapi import ( from .blobapi import BlobAPI from .blob_store import BlobStore -from .client import HttpAPI, HttpError +from .client import HttpAPI, HttpError, ControllerClient 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) diff --git a/ick2/client_tests.py b/ick2/client_tests.py index 1a0bcc5..856f4c3 100644 --- a/ick2/client_tests.py +++ b/ick2/client_tests.py @@ -98,6 +98,73 @@ class HttpAPITests(unittest.TestCase): self.assertEqual(obj, None) +class ControllerClientTests(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) + + self.controller = ick2.ControllerClient() + self.controller.set_http_api(self.client) + self.controller.set_controller_url('https://controller') + self.controller.set_token('SECRET-TOKEN') + self.controller.set_client_name('asterix') + + def test_register_succeeds_on_error(self): + self.session.response = FakeResponse(400) + self.assertEqual(self.controller.register(), None) + + def test_register_succeeds(self): + self.session.response = FakeResponse(200) + self.assertEqual(self.controller.register(), None) + + def test_get_work_raises_exception_on_error(self): + self.session.response = FakeResponse(400) + with self.assertRaises(ick2.HttpError): + self.controller.get_work() + + def test_get_work_succeeds(self): + work = { + 'build_id': 'rome/1' + } + self.session.response = FakeResponse( + 200, body=json.dumps(work), content_type=json_type) + self.assertEqual(self.controller.get_work(), work) + + def test_report_work_raises_exception_on_error(self): + work = { + 'stdout': 'hello, world', + } + self.session.response = FakeResponse(400) + with self.assertRaises(ick2.HttpError): + self.controller.report_work(work) + + def test_report_work_succeeds(self): + work = { + 'stdout': 'hello, world', + } + self.session.response = FakeResponse(200) + self.assertEqual(self.controller.report_work(work), None) + + def test_get_blob_service_url_raises_exception_on_error(self): + self.session.response = FakeResponse(400) + with self.assertRaises(ick2.HttpError): + self.controller.get_blob_service_url() + + def test_get_blob_service_url_succeeds(self): + url = 'https://blobs' + version = { + 'blob_service': url, + } + self.session.response = FakeResponse( + 200, body=json.dumps(version), content_type=json_type) + self.assertEqual(self.controller.get_blob_service_url(), url) + + class FakeHttpSession: def __init__(self): 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) |