diff options
| -rw-r--r-- | cloudinit/sources/DataSourceDigitalOcean.py | 53 | ||||
| -rw-r--r-- | tests/unittests/test_datasource/test_digitalocean.py | 9 |
2 files changed, 30 insertions, 32 deletions
diff --git a/cloudinit/sources/DataSourceDigitalOcean.py b/cloudinit/sources/DataSourceDigitalOcean.py index 985f1663..b7afca93 100644 --- a/cloudinit/sources/DataSourceDigitalOcean.py +++ b/cloudinit/sources/DataSourceDigitalOcean.py @@ -18,11 +18,14 @@ from cloudinit import log as logging from cloudinit import util from cloudinit import sources from cloudinit import url_helper +from cloudinit import ec2_utils +import functools + LOG = logging.getLogger(__name__) BUILTIN_DS_CONFIG = { - 'metadata_url': 'http://169.254.169.254/metadata/v1', + 'metadata_url': 'http://169.254.169.254/metadata/v1/', 'mirrors_url': 'http://mirrors.digitalocean.com/' } MD_RETRIES = 0 @@ -37,9 +40,9 @@ class DataSourceDigitalOcean(sources.DataSource): BUILTIN_DS_CONFIG]) self.metadata_address = self.ds_cfg['metadata_url'] - if self.ds_cfg.get('retries'): + if self.ds_cfg.get('retries'): self.retries = self.ds_cfg['retries'] - else: + else: self.retries = MD_RETRIES if self.ds_cfg.get('timeout'): @@ -48,41 +51,27 @@ class DataSourceDigitalOcean(sources.DataSource): self.timeout = MD_TIMEOUT def get_data(self): - url_map = [ - ('user-data', '/user-data'), - ('vendor-data', '/vendor-data'), - ('public-keys', '/public-keys'), - ('region', '/region'), - ('id', '/id'), - ('hostname', '/hostname'), - ] - - found = False - for (key, path) in url_map: - try: - resp = url_helper.readurl(url=self.metadata_address + path, - timeout=self.timeout, - retries=self.retries) - if resp.code == 200: - found = True - self.metadata[key] = resp.contents - else: - LOG.warn("Path: %s returned %s", path, resp.code) - return False - except url_helper.UrlError as e: - LOG.warn("Path: %s raised exception: %s", path, e) - return False - - return found + caller = functools.partial(util.read_file_or_url, timeout=self.timeout, + retries=self.retries) + md = ec2_utils.MetadataMaterializer(str(caller(self.metadata_address)), + base_url=self.metadata_address, + caller=caller) + + self.metadata = md.materialize() + + if self.metadata.get('id'): + return True + else: + return False def get_userdata_raw(self): - return self.metadata['user-data'] + return "\n".join(self.metadata['user-data']) def get_vendordata_raw(self): - return self.metadata['vendor-data'] + return "\n".join(self.metadata['vendor-data']) def get_public_ssh_keys(self): - return self.metadata['public-keys'].splitlines() + return self.metadata['public-keys'].splitlines() @property def availability_zone(self): diff --git a/tests/unittests/test_datasource/test_digitalocean.py b/tests/unittests/test_datasource/test_digitalocean.py index 9576e042..559a4f9f 100644 --- a/tests/unittests/test_datasource/test_digitalocean.py +++ b/tests/unittests/test_datasource/test_digitalocean.py @@ -26,7 +26,16 @@ from cloudinit.sources import DataSourceDigitalOcean from .. import helpers as test_helpers +# Abbreviated for the test +DO_INDEX = """id + hostname + user-data + vendor-data + public-keys + region""" + DO_META = { + '': DO_INDEX, 'user-data': '#!/bin/bash\necho "user-data"', 'vendor-data': '#!/bin/bash\necho "vendor-data"', 'public-keys': 'ssh-rsa AAAAB3NzaC1yc2EAAAA... neal@digitalocean.com', |