diff --git a/teeth_agent/errors.py b/teeth_agent/errors.py index beeac26fd..c796a56f5 100644 --- a/teeth_agent/errors.py +++ b/teeth_agent/errors.py @@ -61,3 +61,15 @@ class HeartbeatError(errors.RESTError): def __init__(self, details): super(HeartbeatError, self).__init__() self.details = details + + +class ImageDownloadError(errors.RESTError): + """Error raised when an image cannot be downloaded.""" + + message = 'Error downloading image.' + + def __init__(self, image_id, image_url): + details = 'Image with id {} not found at url {}.'.format(image_id, + image_url) + super(ImageDownloadError, self).__init__() + self.details = details diff --git a/teeth_agent/standby.py b/teeth_agent/standby.py index ac93fb23f..871c74cfb 100644 --- a/teeth_agent/standby.py +++ b/teeth_agent/standby.py @@ -48,10 +48,11 @@ def _write_image(image_info, configdrive='configdrive', device='/dev/sda'): def _download_image(image_info): - resp = requests.get(image_info['urls'][0], stream=True) + # TODO(jimrollenhagen) try all URLs + url = image_info['urls'][0] + resp = requests.get(url, stream=True) if resp.status_code != 200: - # TODO(jimrollenhagen) define a better exception - raise Exception + raise errors.ImageDownloadError(image_info['id'], url) image_location = _image_location(image_info) with open(image_location, 'wb') as f: for chunk in resp.iter_content(1024 * 1024): diff --git a/teeth_agent/tests/standby_agent.py b/teeth_agent/tests/standby_agent.py index c5f12b7de..4d0290ab7 100644 --- a/teeth_agent/tests/standby_agent.py +++ b/teeth_agent/tests/standby_agent.py @@ -120,7 +120,9 @@ class TestBaseTeethAgent(unittest.TestCase): image_info = self._build_fake_image_info() response = requests_mock.return_value response.status_code = 404 - self.assertRaises(Exception, standby._download_image, image_info) + self.assertRaises(errors.ImageDownloadError, + standby._download_image, + image_info) @mock.patch('teeth_agent.standby._verify_image', autospec=True) @mock.patch('__builtin__.open', autospec=True)