Merge "Refactor the image download and checksum computation bits"

ironic_python_agent/extensions/standby.py

@@ -115,7 +115,36 @@ def _write_configdrive_to_partition(configdrive, device):
              totaltime))
 
 
-def _request_url(image_info, url):
+class ImageDownload(object):
+    """Helper class that opens a HTTP connection to download an image.
+
+    This class opens a HTTP connection to download an image from a URL
+    and create an iterator so the image can be downloaded in chunks. The
+    MD5 hash of the image being downloaded is calculated on-the-fly.
+    """
+
+    def __init__(self, image_info, time_obj=None):
+        self._md5checksum = hashlib.md5()
+        self._time = time_obj or time.time()
+        self._request = None
+
+        for url in image_info['urls']:
+            try:
+                LOG.info("Attempting to download image from {0}".format(url))
+                self._request = self._download_file(image_info, url)
+            except errors.ImageDownloadError as e:
+                failtime = time.time() - self._time
+                log_msg = ('Image download failed. URL: {0}; time: {1} '
+                           'seconds. Error: {2}')
+                LOG.warning(log_msg.format(url, failtime, e.details))
+                continue
+            else:
+                break
+        else:
+            msg = 'Image download failed for all URLs.'
+            raise errors.ImageDownloadError(image_info['id'], msg)
+
+    def _download_file(self, image_info, url):
         no_proxy = image_info.get('no_proxy')
         if no_proxy:
             os.environ['no_proxy'] = no_proxy
@@ -127,30 +156,34 @@ def _request_url(image_info, url):
             raise errors.ImageDownloadError(image_info['id'], msg)
         return resp
 
+    def __iter__(self):
+        for chunk in self._request.iter_content(IMAGE_CHUNK_SIZE):
+            self._md5checksum.update(chunk)
+            yield chunk
+
+    def md5sum(self):
+        return self._md5checksum.hexdigest()
+
+
+def _verify_image(image_info, image_location, checksum):
+    LOG.debug('Verifying image at {0} against MD5 checksum '
+              '{1}'.format(image_location, checksum))
+    if checksum != image_info['checksum']:
+        LOG.error(errors.ImageChecksumError.details_str.format(
+            image_location, image_info['id'],
+            image_info['checksum'], checksum))
+        raise errors.ImageChecksumError(image_location, image_info['id'],
+                                        image_info['checksum'], checksum)
+
 
 def _download_image(image_info):
     starttime = time.time()
-    resp = None
-    for url in image_info['urls']:
-        try:
-            LOG.info("Attempting to download image from {0}".format(url))
-            resp = _request_url(image_info, url)
-        except errors.ImageDownloadError as e:
-            failtime = time.time() - starttime
-            log_msg = ('Image download failed. URL: {0}; time: {1} seconds. '
-                       'Error: {2}')
-            LOG.warning(log_msg.format(url, failtime, e.details))
-            continue
-        else:
-            break
-    if resp is None:
-        msg = 'Image download failed for all URLs.'
-        raise errors.ImageDownloadError(image_info['id'], msg)
-
     image_location = _image_location(image_info)
+    image_download = ImageDownload(image_info, time_obj=starttime)
+
     with open(image_location, 'wb') as f:
         try:
-            for chunk in resp.iter_content(IMAGE_CHUNK_SIZE):
+            for chunk in image_download:
                 f.write(chunk)
         except Exception as e:
             msg = 'Unable to write image to {0}. Error: {1}'.format(
@@ -160,30 +193,7 @@ def _download_image(image_info):
     totaltime = time.time() - starttime
     LOG.info("Image downloaded from {0} in {1} seconds".format(image_location,
                                                                totaltime))
-
+    _verify_image(image_info, image_location, image_download.md5sum())
-    _verify_image(image_info, image_location)
-
-
-def _verify_image(image_info, image_location):
-    checksum = image_info['checksum']
-    log_msg = 'Verifying image at {0} against MD5 checksum {1}'
-    LOG.debug(log_msg.format(image_location, checksum))
-    hash_ = hashlib.md5()
-    with open(image_location) as image:
-        while True:
-            data = image.read(IMAGE_CHUNK_SIZE)
-            if not data:
-                break
-            hash_.update(data)
-    hash_digest = hash_.hexdigest()
-    if hash_digest == checksum:
-        return True
-
-    LOG.error(errors.ImageChecksumError.details_str.format(
-        image_location, image_info['id'], checksum, hash_digest))
-
-    raise errors.ImageChecksumError(image_location, image_info['id'], checksum,
-                                    hash_digest)
 
 
 def _validate_image_info(ext, image_info=None, **kwargs):

ironic_python_agent/tests/unit/extensions/test_standby.py

@@ -28,12 +28,7 @@ else:
     OPEN_FUNCTION_NAME = 'builtins.open'
 
 
-class TestStandbyExtension(test_base.BaseTestCase):
+def _build_fake_image_info():
-    def setUp(self):
-        super(TestStandbyExtension, self).setUp()
-        self.agent_extension = standby.StandbyExtension()
-
-    def _build_fake_image_info(self):
     return {
         'id': 'fake_id',
         'urls': [
@@ -42,12 +37,18 @@ class TestStandbyExtension(test_base.BaseTestCase):
         'checksum': 'abc123'
     }
 
+
+class TestStandbyExtension(test_base.BaseTestCase):
+    def setUp(self):
+        super(TestStandbyExtension, self).setUp()
+        self.agent_extension = standby.StandbyExtension()
+
     def test_validate_image_info_success(self):
-        standby._validate_image_info(None, self._build_fake_image_info())
+        standby._validate_image_info(None, _build_fake_image_info())
 
     def test_validate_image_info_missing_field(self):
         for field in ['id', 'urls', 'checksum']:
-            invalid_info = self._build_fake_image_info()
+            invalid_info = _build_fake_image_info()
             del invalid_info[field]
 
             self.assertRaises(errors.InvalidCommandParamsError,
@@ -55,7 +56,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                               invalid_info)
 
     def test_validate_image_info_invalid_urls(self):
-        invalid_info = self._build_fake_image_info()
+        invalid_info = _build_fake_image_info()
         invalid_info['urls'] = 'this_is_not_a_list'
 
         self.assertRaises(errors.InvalidCommandParamsError,
@@ -63,7 +64,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                           invalid_info)
 
     def test_validate_image_info_empty_urls(self):
-        invalid_info = self._build_fake_image_info()
+        invalid_info = _build_fake_image_info()
         invalid_info['urls'] = []
 
         self.assertRaises(errors.InvalidCommandParamsError,
@@ -71,7 +72,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                           invalid_info)
 
     def test_validate_image_info_invalid_checksum(self):
-        invalid_info = self._build_fake_image_info()
+        invalid_info = _build_fake_image_info()
         invalid_info['checksum'] = {'not': 'a string'}
 
         self.assertRaises(errors.InvalidCommandParamsError,
@@ -79,7 +80,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                           invalid_info)
 
     def test_validate_image_info_empty_checksum(self):
-        invalid_info = self._build_fake_image_info()
+        invalid_info = _build_fake_image_info()
         invalid_info['checksum'] = ''
 
         self.assertRaises(errors.InvalidCommandParamsError,
@@ -92,14 +93,14 @@ class TestStandbyExtension(test_base.BaseTestCase):
                           image_info={'foo': 'bar'})
 
     def test_image_location(self):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         location = standby._image_location(image_info)
         self.assertEqual(location, '/tmp/fake_id')
 
     @mock.patch(OPEN_FUNCTION_NAME, autospec=True)
     @mock.patch('ironic_python_agent.utils.execute', autospec=True)
     def test_write_image(self, execute_mock, open_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         device = '/dev/sda'
         location = standby._image_location(image_info)
         script = standby._path_to_script('shell/write_image.sh')
@@ -202,7 +203,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
     @mock.patch(OPEN_FUNCTION_NAME)
     @mock.patch('requests.get')
     def test_download_image(self, requests_mock, open_mock, md5_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         response = requests_mock.return_value
         response.status_code = 200
         response.iter_content.return_value = ['some', 'content']
@@ -226,7 +227,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
     @mock.patch.dict(os.environ, {})
     def test_download_image_proxy(
             self, requests_mock, open_mock, md5_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         proxies = {'http': 'http://a.b.com',
                    'https': 'https://secure.a.b.com'}
         no_proxy = '.example.org,.b.com'
@@ -252,57 +253,40 @@ class TestStandbyExtension(test_base.BaseTestCase):
 
     @mock.patch('requests.get', autospec=True)
     def test_download_image_bad_status(self, requests_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         response = requests_mock.return_value
         response.status_code = 404
         self.assertRaises(errors.ImageDownloadError,
                           standby._download_image,
                           image_info)
 
-    @mock.patch('ironic_python_agent.extensions.standby._verify_image',
+    @mock.patch('hashlib.md5', autospec=True)
-                autospec=True)
     @mock.patch(OPEN_FUNCTION_NAME, autospec=True)
     @mock.patch('requests.get', autospec=True)
     def test_download_image_verify_fails(self, requests_mock, open_mock,
-                                         verify_mock):
+                                         md5_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         response = requests_mock.return_value
         response.status_code = 200
-        verify_mock.side_effect = errors.ImageChecksumError(
+        hexdigest_mock = md5_mock.return_value.hexdigest
-            'foo', '/bar/foo', 'incorrect', 'correct')
+        hexdigest_mock.return_value = 'invalid-checksum'
         self.assertRaises(errors.ImageChecksumError,
                           standby._download_image,
                           image_info)
 
-    @mock.patch(OPEN_FUNCTION_NAME)
+    def test_verify_image_success(self):
-    @mock.patch('hashlib.md5')
+        image_info = _build_fake_image_info()
-    def test_verify_image_success(self, md5_mock, open_mock):
-        image_info = self._build_fake_image_info()
-        hexdigest_mock = md5_mock.return_value.hexdigest
-        hexdigest_mock.return_value = image_info['checksum']
-        file_mock = mock.Mock()
-        file_mock.read.return_value = None
-        open_mock.return_value.__enter__.return_value = file_mock
         image_location = '/foo/bar'
+        checksum = image_info['checksum']
+        standby._verify_image(image_info, image_location, checksum)
 
-        verified = standby._verify_image(image_info, image_location)
+    def test_verify_image_failure(self):
-        self.assertTrue(verified)
+        image_info = _build_fake_image_info()
-        self.assertEqual(1, md5_mock.call_count)
-
-    @mock.patch(OPEN_FUNCTION_NAME)
-    @mock.patch('hashlib.md5')
-    def test_verify_image_failure(self, md5_mock, open_mock):
-        image_info = self._build_fake_image_info()
-        md5_mock.return_value.hexdigest.return_value = 'wrong hash'
-        file_mock = mock.Mock()
-        file_mock.read.return_value = None
-        open_mock.return_value.__enter__.return_value = file_mock
         image_location = '/foo/bar'
-
+        checksum = 'invalid-checksum'
         self.assertRaises(errors.ImageChecksumError,
                           standby._verify_image,
-                          image_info, image_location)
+                          image_info, image_location, checksum)
-        self.assertEqual(md5_mock.call_count, 1)
 
     @mock.patch('ironic_python_agent.hardware.dispatch_to_managers',
                 autospec=True)
@@ -312,7 +296,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                 autospec=True)
     def test_cache_image(self, download_mock, write_mock,
                          dispatch_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         download_mock.return_value = None
         write_mock.return_value = None
         dispatch_mock.return_value = 'manager'
@@ -337,7 +321,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                 autospec=True)
     def test_cache_image_force(self, download_mock, write_mock,
                                dispatch_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         self.agent_extension.cached_image_id = image_info['id']
         download_mock.return_value = None
         write_mock.return_value = None
@@ -365,7 +349,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                 autospec=True)
     def test_cache_image_cached(self, download_mock, write_mock,
                                 dispatch_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         self.agent_extension.cached_image_id = image_info['id']
         download_mock.return_value = None
         write_mock.return_value = None
@@ -400,7 +384,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                            write_mock,
                            dispatch_mock,
                            configdrive_copy_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         location_mock.return_value = '/tmp/configdrive'
         download_mock.return_value = None
         write_mock.return_value = None
@@ -460,7 +444,7 @@ class TestStandbyExtension(test_base.BaseTestCase):
                                           write_mock,
                                           dispatch_mock,
                                           configdrive_copy_mock):
-        image_info = self._build_fake_image_info()
+        image_info = _build_fake_image_info()
         download_mock.return_value = None
         write_mock.return_value = None
         dispatch_mock.return_value = 'manager'
@@ -530,3 +514,23 @@ class TestStandbyExtension(test_base.BaseTestCase):
 
         execute_mock.assert_called_once_with(*command, check_exit_code=[0])
         self.assertEqual('FAILED', failed_result.command_status)
+
+
+class TestImageDownload(test_base.BaseTestCase):
+
+    @mock.patch('hashlib.md5', autospec=True)
+    @mock.patch('requests.get', autospec=True)
+    def test_download_image(self, requests_mock, md5_mock):
+        content = ['SpongeBob', 'SquarePants']
+        response = requests_mock.return_value
+        response.status_code = 200
+        response.iter_content.return_value = content
+
+        image_info = _build_fake_image_info()
+        md5_mock.return_value.hexdigest.return_value = image_info['checksum']
+        image_download = standby.ImageDownload(image_info)
+
+        self.assertEqual(content, list(image_download))
+        requests_mock.assert_called_once_with(image_info['urls'][0],
+                                              stream=True, proxies={})
+        self.assertEqual(image_info['checksum'], image_download.md5sum())