Fix Requests breaking download progress bar
The move to the requests library (dbb242b
) broke the progress bar during
downloading an image with --progress enabled, due to requests returning
a generator, which has no __len__ function. This patch adds an iterable
wrapper which provides the length, sourced from Content-Length headers.
Closes-Bug: #1384664
Change-Id: I48598a824396bc6e7cba69eb3ce32e7c8f30a18a
This commit is contained in:
@@ -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
|
||||
|
@@ -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.
|
||||
|
@@ -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):
|
||||
"""
|
||||
|
@@ -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:
|
||||
|
Reference in New Issue
Block a user