async def test_multiple_files(self):
file1 = ('file1.txt', streams.StringStream('[File One]'))
file2 = ('file2.txt', streams.StringStream('[File Two]'))
file3 = ('file3.txt', streams.StringStream('[File Three]'))
files = AsyncIterator([file1, file2, file3])
stream = streams.ZipStreamReader(files)
data = await stream.read()
zip = zipfile.ZipFile(io.BytesIO(data))
# Verify CRCs
assert zip.testzip() is None
# Check content of included files
zipped1 = zip.open('file1.txt')
assert zipped1.read() == b'[File One]'
zipped2 = zip.open('file2.txt')
assert zipped2.read() == b'[File Two]'
zipped3 = zip.open('file3.txt')
assert zipped3.read() == b'[File Three]'
def download(self, path, revision=None, **kwargs):
if revision and not revision.endswith(settings.DRIVE_IGNORE_VERSION):
# Must make additional request to look up download URL for revision
response = yield from self.make_request(
'GET',
self.build_url('files',
path.identifier,
'revisions',
revision,
alt='json'),
expects=(200, ),
throws=exceptions.MetadataError,
)
data = yield from response.json()
else:
data = yield from self.metadata(path, raw=True)
download_resp = yield from self.make_request(
'GET',
data.get('downloadUrl')
or drive_utils.get_export_link(data['exportLinks']),
expects=(200, ),
throws=exceptions.DownloadError,
)
if 'fileSize' in data:
return streams.ResponseStreamReader(download_resp,
size=data['fileSize'])
# google docs, not drive files, have no way to get the file size
# must buffer the entire file into memory
stream = streams.StringStream((yield from download_resp.read()))
if download_resp.headers.get('Content-Type'):
stream.content_type = download_resp.headers['Content-Type']
return stream
async def test_upload_parts(self, provider, root_provider_fixtures):
responses = [{
'body':
json.dumps(root_provider_fixtures['upload_part_one']),
'status':
201
}, {
'body':
json.dumps(root_provider_fixtures['upload_part_two']),
'status':
201
}]
session_url = 'https://upload.box.com/api/2.0/files/upload_sessions/fake_session_id'
aiohttpretty.register_json_uri('PUT',
session_url,
status=HTTPStatus.CREATED,
responses=responses)
session_metadata = root_provider_fixtures['create_session_metadata']
stream = streams.StringStream('tenbytestr'.encode() * 2)
parts_metadata = await provider._upload_parts(stream, session_metadata)
expected_response = [{
'offset': 10,
'part_id': '37B0FB1B',
'sha1': '3ff00d99585b8da363f9f9955e791ed763e111c1',
'size': 10
}, {
'offset': 20,
'part_id': '1872DEDA',
'sha1': '0ae5fc290c5c5414cdda245ab712a8440376284a',
'size': 10
}]
assert parts_metadata == expected_response
assert len(aiohttpretty.calls) == 2
for call in aiohttpretty.calls:
assert call['method'] == 'PUT'
assert call['uri'] == session_url
call_one = aiohttpretty.calls[0]
assert call_one['headers'] == {
'Authorization': 'Bearer wrote harry potter',
'Content-Length': '10',
'Content-Range': 'bytes 0-9/20',
'Content-Type:': 'application/octet-stream',
'Digest': 'sha={}'.format('pz4mZbOEOesBeUhR1THUF1Oq1bI=')
}
call_two = aiohttpretty.calls[1]
assert call_two['headers'] == {
'Authorization': 'Bearer wrote harry potter',
'Content-Length': '10',
'Content-Range': 'bytes 10-19/20',
'Content-Type:': 'application/octet-stream',
'Digest': 'sha={}'.format('pz4mZbOEOesBeUhR1THUF1Oq1bI=')
}
async def download(self, path, revision=None, range=None, **kwargs):
if revision and not revision.endswith(settings.DRIVE_IGNORE_VERSION):
metadata = await self.metadata(path, revision=revision)
else:
metadata = await self.metadata(path)
download_resp = await self.make_request(
'GET',
metadata.raw.get('downloadUrl')
or drive_utils.get_export_link(metadata.raw),
range=range,
expects=(200, 206),
throws=exceptions.DownloadError,
)
if metadata.size is not None:
return streams.ResponseStreamReader(download_resp,
size=metadata.size)
# google docs, not drive files, have no way to get the file size
# must buffer the entire file into memory
stream = streams.StringStream(await download_resp.read())
if download_resp.headers.get('Content-Type'):
stream.content_type = download_resp.headers['Content-Type']
stream.name = metadata.export_name
return stream
def test_chunking(self):
for chunk_size in range(1, 10):
data = b'the ode to carp'
expected = streams.StringStream(base64.b64encode(data))
stream = streams.Base64EncodeStream(streams.StringStream(data))
hoped = yield from expected.read(chunk_size)
while hoped:
actual = yield from stream.read(chunk_size)
assert actual == hoped
hoped = yield from expected.read(chunk_size)
left_overs = yield from stream.read()
assert left_overs == b''
async def test_1_at_a_time_many_stream(self, blob):
count = 4
stream = streams.MultiStream(
*[streams.StringStream(blob) for _ in range(count)])
for _ in range(count):
for i in range(len(blob)):
assert blob[i:i + 1] == (await stream.read(1))
def test_read(self):
data = b'this is a test'
expected = base64.b64encode(data)
stream = streams.Base64EncodeStream(streams.StringStream(data))
actual = yield from stream.read()
assert expected == actual
def test_other_streams(self):
stream = streams.JSONStream({
'justAStream': streams.StringStream('These are some words')
})
read = yield from stream.read()
assert json.loads(read.decode('utf-8')) == {
'justAStream': 'These are some words'
}
async def download(
self, # type: ignore
path: GitLabPath,
**kwargs):
"""Return a stream to the specified file on GitLab.
There is an endpoint for downloading the raw file directly, but we cannot use it because
GitLab requires periods in the file path to be encoded. Python and aiohttp make this
difficult, though their behavior is arguably correct. See
https://gitlab.com/gitlab-org/gitlab-ce/issues/31470 for details. (Update: this is due to
be fixed in the GL 10.0 release)
API docs: https://docs.gitlab.com/ce/api/repository_files.html#get-file-from-repository
This uses the same endpoint as `_fetch_file_contents`, but relies on the response headers,
which are not returned by that method. It may also be replaced when the above bug is
fixed.
:param str path: The path to the file on GitLab
:param dict \*\*kwargs: Ignored
:raises: :class:`waterbutler.core.exceptions.DownloadError`
"""
url = self._build_file_url(path)
resp = await self.make_request(
'GET',
url,
expects=(200, ),
throws=exceptions.DownloadError,
)
raw_data = (await resp.read()).decode("utf-8")
data = None
try:
data = json.loads(raw_data)
except json.decoder.JSONDecodeError:
# GitLab API sometimes returns ruby hashes instead of json
# see: https://gitlab.com/gitlab-org/gitlab-ce/issues/31790
# fixed in GL v9.5
data = self._convert_ruby_hash_to_dict(raw_data)
raw = base64.b64decode(data['content'])
mdict_options = {}
mimetype = mimetypes.guess_type(path.full_path)[0]
if mimetype is not None:
mdict_options['CONTENT-TYPE'] = mimetype
mdict = aiohttp.multidict.MultiDict(resp.headers)
mdict.update(mdict_options)
resp.headers = mdict
resp.content = streams.StringStream(raw)
return streams.ResponseStreamReader(resp, len(raw))
def test_github_at_once(self):
stream = streams.JSONStream({
'encoding': 'base64',
'content': streams.Base64EncodeStream(streams.StringStream('These are some words')),
})
buffer = yield from stream.read()
assert json.loads(buffer.decode('utf-8')) == {
'encoding': 'base64',
'content': 'VGhlc2UgYXJlIHNvbWUgd29yZHM='
}
def test_download_stream(self):
data = b'freddie brian john roger'
stream = streams.StringStream(data)
stream.content_type = 'application/octet-stream'
self.mock_provider.download = utils.MockCoroutine(return_value=stream)
resp = yield self.http_client.fetch(
self.get_url('/file?provider=queenhub&path=/freddie.png'), )
assert resp.body == data
calls = self.mock_provider.download.call_args_list
assert len(calls) == 1
args, kwargs = calls[0]
assert kwargs.get('action') == 'download'
def test_upload_create(self, provider):
file_name = 'upload.txt'
file_folder = '/'
file_path = os.path.join(file_folder, file_name)
file_content = b'Test Upload Content'
file_stream = streams.StringStream(file_content)
path = yield from provider.validate_path(file_path)
metadata, created = yield from provider.upload(file_stream, path)
assert metadata['name'] == file_name
assert metadata['path'] == file_path
assert metadata['size'] == len(file_content)
assert created is True
async def test_upload_create(self, provider):
file_name = 'upload.txt'
file_folder = '/'
file_path = os.path.join(file_folder, file_name)
file_content = b'Test Upload Content'
file_stream = streams.StringStream(file_content)
path = await provider.validate_path(file_path)
metadata, created = await provider.upload(file_stream, path)
assert metadata.name == file_name
assert metadata.path == file_path
assert metadata.size == len(file_content)
assert created is True
def test_other_streams_1_at_a_time(self):
stream = streams.JSONStream({
'justAStream': streams.StringStream('These are some words')
})
buffer = b''
chunk = yield from stream.read(1)
while chunk:
buffer += chunk
chunk = yield from stream.read(1)
assert json.loads(buffer.decode('utf-8')) == {
'justAStream': 'These are some words'
}
def test_single_file(self):
file = ('filename.extension', streams.StringStream('[File Content]'))
stream = streams.ZipStreamReader(file)
data = yield from stream.read()
zip = zipfile.ZipFile(io.BytesIO(data))
# Verify CRCs
assert zip.testzip() is None
result = zip.open('filename.extension')
# Check content of included file
assert result.read() == b'[File Content]'
def test_download_content_type_does_not_switch(self):
"""mime_types should not override file extension not in the dict
"""
data = b'freddie brian john roger'
stream = streams.StringStream(data)
stream.content_type = 'application/octet-stream'
self.mock_provider.download = utils.MockCoroutine(return_value=stream)
resp = yield self.http_client.fetch(
self.get_url('/file?provider=queenhub&path=/freddie.png'), )
assert resp.body == data
assert resp.headers['Content-Type'] == 'application/octet-stream'
calls = self.mock_provider.download.call_args_list
assert len(calls) == 1
args, kwargs = calls[0]
assert kwargs.get('action') == 'download'
async def test_multi_chunk(self, blob):
stream = streams.StringStream(blob)
cutoff_stream_one = streams.CutoffStream(stream, 10)
data_one = await cutoff_stream_one.read()
assert len(data_one) == 10
assert data_one == blob[0:10]
cutoff_stream_two = streams.CutoffStream(stream, 10)
data_two = await cutoff_stream_two.read()
assert len(data_two) == 10
assert data_two == blob[10:20]
remainder = await stream.read()
assert len(remainder) == 30
assert remainder == blob[20:50]
async def test_file(self):
stream = streams.FormDataStream(
file=streams.StringStream('Empty chairs at empty tables'))
data = await stream.read()
expected = sorted([
'--thisisaknownvalue', 'Content-Disposition: file; name="file"',
'Content-Type: application/octet-stream',
'Content-Transfer-Encoding: binary', '',
'Empty chairs at empty tables', '--thisisaknownvalue--', ''
])
actual = sorted(data.decode('utf-8').split('\r\n'))
assert expected == actual
async def download(
self, # type: ignore
path: GoogleDrivePath,
revision: str = None,
range: Tuple[int, int] = None,
**kwargs) -> streams.BaseStream:
"""Download the file at `path`. If `revision` is present, attempt to download that revision
of the file. See **Revisions** in the class doctring for an explanation of this provider's
revision handling. The actual revision handling is done in `_file_metadata()`.
Quirks:
Google docs don't have a size until they're exported, so WB must download them, then
re-stream them as a StringStream.
:param GoogleDrivePath path: the file to download
:param str revision: the id of a particular version to download
:param tuple(int, int) range: range of bytes to download in this request
:rtype: streams.ResponseStreamReader
:rtype: streams.StringStream
:returns: For GDocs, a StringStream. All others, a ResponseStreamReader.
"""
metadata = await self.metadata(path, revision=revision)
download_resp = await self.make_request(
'GET',
metadata.raw.get('downloadUrl')
or utils.get_export_link(metadata.raw), # type: ignore
range=range,
expects=(200, 206),
throws=exceptions.DownloadError,
)
if metadata.size is not None: # type: ignore
return streams.ResponseStreamReader(
download_resp, size=metadata.size_as_int) # type: ignore
# google docs, not drive files, have no way to get the file size
# must buffer the entire file into memory
stream = streams.StringStream(await download_resp.read())
if download_resp.headers.get('Content-Type'):
# TODO: Add these properties to base class officially, instead of as one-off
stream.content_type = download_resp.headers[
'Content-Type'] # type: ignore
stream.name = metadata.export_name # type: ignore
return stream
async def test_github(self):
stream = streams.JSONStream({
'encoding': 'base64',
'content': streams.Base64EncodeStream(streams.StringStream('These are some words')),
})
buffer = b''
chunk = await stream.read(1)
while chunk:
buffer += chunk
chunk = await stream.read(1)
assert json.loads(buffer.decode('utf-8')) == {
'encoding': 'base64',
'content': 'VGhlc2UgYXJlIHNvbWUgd29yZHM='
}
def test_download_stream(self):
data = b'freddie brian john roger'
stream = streams.StringStream(data)
stream.content_type = 'application/octet-stream'
zipstream = streams.ZipStreamReader(('file.txt', stream))
self.mock_provider.zip = utils.MockCoroutine(return_value=zipstream)
resp = yield self.http_client.fetch(
self.get_url('/zip?provider=queenhub&path=/freddie.png'), )
zip = zipfile.ZipFile(io.BytesIO(resp.body))
assert zip.testzip() is None
assert zip.open('file.txt').read() == data
def test_download_content_type_switches(self):
"""waterbutler.core.mime_types contains content type
overrides.
"""
data = b'freddie brian john roger'
stream = streams.StringStream(data)
stream.content_type = 'application/octet-stream'
self.mock_provider.download = utils.MockCoroutine(return_value=stream)
resp = yield self.http_client.fetch(
self.get_url('/file?provider=queenhub&path=/freddie.md'), )
assert resp.body == data
assert resp.headers['Content-Type'] == 'text/x-markdown'
calls = self.mock_provider.download.call_args_list
assert len(calls) == 1
args, kwargs = calls[0]
assert kwargs.get('action') == 'download'
def test_download_stream_range(self):
data = b'freddie brian john roger'
stream = streams.StringStream(data)
stream.name = 'foo'
stream.partial = True
stream.content_range = '0-{}/{}'.format(len(data) - 1, len(data))
stream.content_type = 'application/octet-stream'
self.mock_provider.download = utils.MockCoroutine(return_value=stream)
resp = yield self.http_client.fetch(
self.get_url('/file?provider=queenhub&path=/freddie.png'),
headers={'Range': 'bytes=0-'})
assert resp.code == 206
assert resp.body == data
calls = self.mock_provider.download.call_args_list
assert len(calls) == 1
args, kwargs = calls[0]
assert kwargs.get('action') == 'download'
assert kwargs.get('range') == (0, None)
async def test_subchunk(self, blob):
stream = streams.StringStream(blob)
cutoff_stream = streams.CutoffStream(stream, 20)
subchunk_one = await cutoff_stream.read(7)
assert len(subchunk_one) == 7
assert subchunk_one == blob[0:7]
subchunk_two = await cutoff_stream.read(7)
assert len(subchunk_two) == 7
assert subchunk_two == blob[7:14]
subchunk_three = await cutoff_stream.read(7)
assert len(subchunk_three) == 6
assert subchunk_three == blob[14:20]
subchunk_four = await cutoff_stream.read(7)
assert len(subchunk_four) == 0
assert subchunk_four == b''
remainder = await stream.read()
assert len(remainder) == 30
assert remainder == blob[20:50]
async def test_1_at_a_time_single_stream(self, blob):
stream = streams.MultiStream(streams.StringStream(blob))
for i in range(len(blob)):
assert blob[i:i + 1] == (await stream.read(1))
async def test_double_same_stream(self, blob):
stream = streams.MultiStream(streams.StringStream(blob),
streams.StringStream(blob))
data = await stream.read()
assert data == (blob * 2)
async def test_single_stream(self, blob):
stream = streams.MultiStream(streams.StringStream(blob))
data = await stream.read()
assert data == blob
async def test_works(self):
data = b'This here be a string yar'
stream = streams.StringStream(data)
read = await stream.read()
assert data == read
def test_no_cutoff_exception(self, blob):
stream = streams.StringStream(blob)
with pytest.raises(TypeError):
streams.CutoffStream(stream)
async def test_one_chunk(self, blob):
stream = streams.StringStream(blob)
cutoff_stream = streams.CutoffStream(stream, len(blob))
data = await cutoff_stream.read()
assert len(data) == len(blob)
assert data == blob
