Skip to content

Commit

Permalink
Fix Requests breaking download progress bar
Browse files Browse the repository at this point in the history
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
  • Loading branch information
kragniz committed Jan 4, 2015
1 parent 7ce1ff9 commit df02ee8
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 5 deletions.
15 changes: 15 additions & 0 deletions glanceclient/common/utils.py
Expand Up @@ -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
5 changes: 3 additions & 2 deletions glanceclient/v1/images.py
Expand Up @@ -140,14 +140,15 @@ def data(self, image, do_checksum=True, **kwargs):
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.
Expand Down
8 changes: 5 additions & 3 deletions glanceclient/v2/images.py
Expand Up @@ -115,10 +115,12 @@ def data(self, image_id, do_checksum=True):
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):
"""
Expand Down
13 changes: 13 additions & 0 deletions tests/v2/test_shell_v2.py
Expand Up @@ -341,6 +341,19 @@ def test_image_upload(self):
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:
Expand Down

0 comments on commit df02ee8

Please sign in to comment.