diff --git a/glanceclient/common/utils.py b/glanceclient/common/utils.py index fe7a51a5..61c526ee 100644 --- a/glanceclient/common/utils.py +++ b/glanceclient/common/utils.py @@ -421,3 +421,18 @@ def safe_header(name, value): return name, "{SHA1}%s" % d else: return name, value + + +class IterableWithLength(object): + def __init__(self, iterable, length): + self.iterable = iterable + self.length = length + + def __iter__(self): + return self.iterable + + def next(self): + return next(self.iterable) + + def __len__(self): + return self.length diff --git a/glanceclient/v1/images.py b/glanceclient/v1/images.py index 87060c23..e3a6e695 100644 --- a/glanceclient/v1/images.py +++ b/glanceclient/v1/images.py @@ -140,14 +140,15 @@ class ImageManager(base.ManagerWithFind): image_id = base.getid(image) resp, body = self.client.get('/v1/images/%s' % urlparse.quote(str(image_id))) + content_length = int(resp.headers.get('content-length', 0)) checksum = resp.headers.get('x-image-meta-checksum', None) if do_checksum and checksum is not None: - return utils.integrity_iter(body, checksum) + body = utils.integrity_iter(body, checksum) return_request_id = kwargs.get('return_req_id', None) if return_request_id is not None: return_request_id.append(resp.headers.get(OS_REQ_ID_HDR, None)) - return body + return utils.IterableWithLength(body, content_length) def list(self, **kwargs): """Get a list of images. diff --git a/glanceclient/v2/images.py b/glanceclient/v2/images.py index 6ec92507..b49c5ee0 100644 --- a/glanceclient/v2/images.py +++ b/glanceclient/v2/images.py @@ -115,10 +115,12 @@ class Controller(object): url = '/v2/images/%s/file' % image_id resp, body = self.http_client.get(url) checksum = resp.headers.get('content-md5', None) + content_length = int(resp.headers.get('content-length', 0)) + if do_checksum and checksum is not None: - return utils.integrity_iter(body, checksum) - else: - return body + body = utils.integrity_iter(body, checksum) + + return utils.IterableWithLength(body, content_length) def upload(self, image_id, image_data, image_size=None): """ diff --git a/tests/v2/test_shell_v2.py b/tests/v2/test_shell_v2.py index 23894eab..3480e5c9 100644 --- a/tests/v2/test_shell_v2.py +++ b/tests/v2/test_shell_v2.py @@ -341,6 +341,19 @@ class ShellV2Test(testtools.TestCase): test_shell.do_image_upload(self.gc, args) mocked_upload.assert_called_once_with('IMG-01', 'testfile', 1024) + def test_image_download(self): + args = self._make_args( + {'id': 'IMG-01', 'file': 'test', 'progress': True}) + + with mock.patch.object(self.gc.images, 'data') as mocked_data: + def _data(): + for c in 'abcedf': + yield c + mocked_data.return_value = utils.IterableWithLength(_data(), 5) + + test_shell.do_image_download(self.gc, args) + mocked_data.assert_called_once_with('IMG-01') + def test_do_image_delete(self): args = self._make_args({'id': 'pass', 'file': 'test'}) with mock.patch.object(self.gc.images, 'delete') as mocked_delete: