# Copyright 2022 Google LLC

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import pytest

from google.cloud.storage import Blob

from google.cloud.storage import Client

from google.cloud.storage import transfer_manager

from google.cloud.storage.retry import DEFAULT_RETRY

from google.api_core import exceptions

from google.cloud.storage.exceptions import DataCorruption

import os

import tempfile

import mock

import pickle

BLOB_TOKEN_STRING = "blob token"

FAKE_CONTENT_TYPE = "text/fake"

UPLOAD_KWARGS = {"content-type": FAKE_CONTENT_TYPE}

FAKE_RESULT = "nothing to see here"

FAKE_ENCODING = "fake_gzip"

DOWNLOAD_KWARGS = {"accept-encoding": FAKE_ENCODING}

CHUNK_SIZE = 8

HOSTNAME = "https://example.com"

URL = "https://example.com/bucket/blob"

USER_AGENT = "agent"

EXPECTED_UPLOAD_KWARGS = {

"command": "tm.upload_many",

**UPLOAD_KWARGS,

}

EXPECTED_DOWNLOAD_KWARGS = {

"command": "tm.download_many",

**DOWNLOAD_KWARGS,

}

# Used in subprocesses only, so excluded from coverage

def _validate_blob_token_in_subprocess(

maybe_pickled_blob, method_name, path_or_file, **kwargs

): # pragma: NO COVER

assert pickle.loads(maybe_pickled_blob) == BLOB_TOKEN_STRING

assert "filename" in method_name

assert path_or_file.startswith("file")

assert kwargs == EXPECTED_UPLOAD_KWARGS or kwargs == EXPECTED_DOWNLOAD_KWARGS

return FAKE_RESULT

def test_upload_many_with_filenames():

FILE_BLOB_PAIRS = [

("file_a.txt", mock.Mock(spec=Blob)),

("file_b.txt", mock.Mock(spec=Blob)),

]

expected_upload_kwargs = EXPECTED_UPLOAD_KWARGS.copy()

expected_upload_kwargs["if_generation_match"] = 0

for _, blob_mock in FILE_BLOB_PAIRS:

blob_mock._handle_filename_and_upload.return_value = FAKE_RESULT

results = transfer_manager.upload_many(

FILE_BLOB_PAIRS,

skip_if_exists=True,

upload_kwargs=UPLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

)

for filename, mock_blob in FILE_BLOB_PAIRS:

mock_blob._handle_filename_and_upload.assert_any_call(

filename, **expected_upload_kwargs

)

for result in results:

assert result == FAKE_RESULT

def test_upload_many_with_file_objs():

FILE_BLOB_PAIRS = [

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

]

expected_upload_kwargs = EXPECTED_UPLOAD_KWARGS.copy()

expected_upload_kwargs["if_generation_match"] = 0

for _, blob_mock in FILE_BLOB_PAIRS:

blob_mock._prep_and_do_upload.return_value = FAKE_RESULT

results = transfer_manager.upload_many(

FILE_BLOB_PAIRS,

skip_if_exists=True,

upload_kwargs=UPLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

)

for file, mock_blob in FILE_BLOB_PAIRS:

mock_blob._prep_and_do_upload.assert_any_call(file, **expected_upload_kwargs)

for result in results:

assert result == FAKE_RESULT

def test_upload_many_passes_concurrency_options():

FILE_BLOB_PAIRS = [

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

]

MAX_WORKERS = 7

DEADLINE = 10

with mock.patch("concurrent.futures.ThreadPoolExecutor") as pool_patch, mock.patch(

"concurrent.futures.wait"

) as wait_patch:

transfer_manager.upload_many(

FILE_BLOB_PAIRS,

deadline=DEADLINE,

worker_type=transfer_manager.THREAD,

max_workers=MAX_WORKERS,

)

pool_patch.assert_called_with(max_workers=MAX_WORKERS)

wait_patch.assert_called_with(mock.ANY, timeout=DEADLINE, return_when=mock.ANY)

def test_threads_deprecation_with_upload():

FILE_BLOB_PAIRS = [

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

]

MAX_WORKERS = 7

DEADLINE = 10

with mock.patch("concurrent.futures.ThreadPoolExecutor") as pool_patch, mock.patch(

"concurrent.futures.wait"

) as wait_patch:

with pytest.warns():

transfer_manager.upload_many(

FILE_BLOB_PAIRS, deadline=DEADLINE, threads=MAX_WORKERS

)

pool_patch.assert_called_with(max_workers=MAX_WORKERS)

wait_patch.assert_called_with(mock.ANY, timeout=DEADLINE, return_when=mock.ANY)

def test_threads_deprecation_conflict_with_upload():

FILE_BLOB_PAIRS = [

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

(tempfile.TemporaryFile(), mock.Mock(spec=Blob)),

]

MAX_WORKERS = 7

DEADLINE = 10

with pytest.raises(ValueError):

transfer_manager.upload_many(

FILE_BLOB_PAIRS,

deadline=DEADLINE,

threads=5,

worker_type=transfer_manager.THREAD,

max_workers=MAX_WORKERS,

)

def test_upload_many_suppresses_exceptions():

FILE_BLOB_PAIRS = [

("file_a.txt", mock.Mock(spec=Blob)),

("file_b.txt", mock.Mock(spec=Blob)),

]

for _, mock_blob in FILE_BLOB_PAIRS:

mock_blob._handle_filename_and_upload.side_effect = ConnectionError()

results = transfer_manager.upload_many(

FILE_BLOB_PAIRS, worker_type=transfer_manager.THREAD

)

for result in results:

assert isinstance(result, ConnectionError)

def test_upload_many_raises_exceptions():

FILE_BLOB_PAIRS = [

("file_a.txt", mock.Mock(spec=Blob)),

("file_b.txt", mock.Mock(spec=Blob)),

]

for _, mock_blob in FILE_BLOB_PAIRS:

mock_blob._handle_filename_and_upload.side_effect = ConnectionError()

with pytest.raises(ConnectionError):

transfer_manager.upload_many(

FILE_BLOB_PAIRS, raise_exception=True, worker_type=transfer_manager.THREAD

)

def test_upload_many_suppresses_412_with_skip_if_exists():

FILE_BLOB_PAIRS = [

("file_a.txt", mock.Mock(spec=Blob)),

("file_b.txt", mock.Mock(spec=Blob)),

]

for _, mock_blob in FILE_BLOB_PAIRS:

mock_blob._handle_filename_and_upload.side_effect = (

exceptions.PreconditionFailed("412")

)

results = transfer_manager.upload_many(

FILE_BLOB_PAIRS,

skip_if_exists=True,

raise_exception=True,

worker_type=transfer_manager.THREAD,

)

for result in results:

assert isinstance(result, exceptions.PreconditionFailed)

def test_upload_many_with_processes():

# Mocks are not pickleable, so we send token strings over the wire.

FILE_BLOB_PAIRS = [

("file_a.txt", BLOB_TOKEN_STRING),

("file_b.txt", BLOB_TOKEN_STRING),

]

with mock.patch(

"google.cloud.storage.transfer_manager._call_method_on_maybe_pickled_blob",

new=_validate_blob_token_in_subprocess,

):

results = transfer_manager.upload_many(

FILE_BLOB_PAIRS,

upload_kwargs=UPLOAD_KWARGS,

worker_type=transfer_manager.PROCESS,

raise_exception=True,

)

for result in results:

assert result == FAKE_RESULT

def test_upload_many_with_processes_rejects_file_obj():

# Mocks are not pickleable, so we send token strings over the wire.

FILE_BLOB_PAIRS = [

("file_a.txt", BLOB_TOKEN_STRING),

(tempfile.TemporaryFile(), BLOB_TOKEN_STRING),

]

with mock.patch(

"google.cloud.storage.transfer_manager._call_method_on_maybe_pickled_blob",

new=_validate_blob_token_in_subprocess,

):

with pytest.raises(ValueError):

transfer_manager.upload_many(

FILE_BLOB_PAIRS,

upload_kwargs=UPLOAD_KWARGS,

worker_type=transfer_manager.PROCESS,

)

def test_download_many_with_filenames():

BLOB_FILE_PAIRS = [

(mock.Mock(spec=Blob), "file_a.txt"),

(mock.Mock(spec=Blob), "file_b.txt"),

]

for blob_mock, _ in BLOB_FILE_PAIRS:

blob_mock._handle_filename_and_download.return_value = FAKE_RESULT

results = transfer_manager.download_many(

BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

)

for mock_blob, file in BLOB_FILE_PAIRS:

mock_blob._handle_filename_and_download.assert_any_call(

file, **EXPECTED_DOWNLOAD_KWARGS

)

for result in results:

assert result == FAKE_RESULT

def test_download_many_with_skip_if_exists():

with tempfile.NamedTemporaryFile() as tf:

BLOB_FILE_PAIRS = [

(mock.Mock(spec=Blob), "file_a.txt"),

(mock.Mock(spec=Blob), tf.name),

]

for blob_mock, _ in BLOB_FILE_PAIRS:

blob_mock._handle_filename_and_download.return_value = FAKE_RESULT

results = transfer_manager.download_many(

BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

skip_if_exists=True,

)

mock_blob, file = BLOB_FILE_PAIRS[0]

mock_blob._handle_filename_and_download.assert_any_call(

file, **EXPECTED_DOWNLOAD_KWARGS

)

mock_blob, _ = BLOB_FILE_PAIRS[1]

mock_blob._handle_filename_and_download.assert_not_called()

for result in results:

assert result == FAKE_RESULT

def test_download_many_with_file_objs():

BLOB_FILE_PAIRS = [

(mock.Mock(spec=Blob), tempfile.TemporaryFile()),

(mock.Mock(spec=Blob), tempfile.TemporaryFile()),

]

for blob_mock, _ in BLOB_FILE_PAIRS:

blob_mock._prep_and_do_download.return_value = FAKE_RESULT

results = transfer_manager.download_many(

BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

)

for mock_blob, file in BLOB_FILE_PAIRS:

mock_blob._prep_and_do_download.assert_any_call(file, **DOWNLOAD_KWARGS)

for result in results:

assert result == FAKE_RESULT

def test_download_many_passes_concurrency_options():

BLOB_FILE_PAIRS = [

(mock.Mock(spec=Blob), tempfile.TemporaryFile()),

(mock.Mock(spec=Blob), tempfile.TemporaryFile()),

]

MAX_WORKERS = 7

DEADLINE = 10

with mock.patch("concurrent.futures.ThreadPoolExecutor") as pool_patch, mock.patch(

"concurrent.futures.wait"

) as wait_patch:

transfer_manager.download_many(

BLOB_FILE_PAIRS,

deadline=DEADLINE,

worker_type=transfer_manager.THREAD,

max_workers=MAX_WORKERS,

)

pool_patch.assert_called_with(max_workers=MAX_WORKERS)

wait_patch.assert_called_with(mock.ANY, timeout=DEADLINE, return_when=mock.ANY)

def test_download_many_suppresses_exceptions():

BLOB_FILE_PAIRS = [

(mock.Mock(spec=Blob), "file_a.txt"),

(mock.Mock(spec=Blob), "file_b.txt"),

]

for mock_blob, _ in BLOB_FILE_PAIRS:

mock_blob._handle_filename_and_download.side_effect = ConnectionError()

results = transfer_manager.download_many(

BLOB_FILE_PAIRS, worker_type=transfer_manager.THREAD

)

for result in results:

assert isinstance(result, ConnectionError)

def test_download_many_raises_exceptions():

BLOB_FILE_PAIRS = [

(mock.Mock(spec=Blob), "file_a.txt"),

(mock.Mock(spec=Blob), "file_b.txt"),

]

for mock_blob, _ in BLOB_FILE_PAIRS:

mock_blob._handle_filename_and_download.side_effect = ConnectionError()

with pytest.raises(ConnectionError):

transfer_manager.download_many(

BLOB_FILE_PAIRS, raise_exception=True, worker_type=transfer_manager.THREAD

)

def test_download_many_with_processes():

# Mocks are not pickleable, so we send token strings over the wire.

BLOB_FILE_PAIRS = [

(BLOB_TOKEN_STRING, "file_a.txt"),

(BLOB_TOKEN_STRING, "file_b.txt"),

]

with mock.patch(

"google.cloud.storage.transfer_manager._call_method_on_maybe_pickled_blob",

new=_validate_blob_token_in_subprocess,

):

results = transfer_manager.download_many(

BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.PROCESS,

)

for result in results:

assert result == FAKE_RESULT

def test_download_many_with_processes_rejects_file_obj():

# Mocks are not pickleable, so we send token strings over the wire.

BLOB_FILE_PAIRS = [

(BLOB_TOKEN_STRING, "file_a.txt"),

(BLOB_TOKEN_STRING, tempfile.TemporaryFile()),

]

with mock.patch(

"google.cloud.storage.transfer_manager._call_method_on_maybe_pickled_blob",

new=_validate_blob_token_in_subprocess,

):

with pytest.raises(ValueError):

transfer_manager.download_many(

BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.PROCESS,

)

def test_upload_many_from_filenames():

bucket = mock.Mock()

FILENAMES = ["file_a.txt", "file_b.txt"]

ROOT = "mypath/"

PREFIX = "myprefix/"

KEY_NAME = "keyname"

BLOB_CONSTRUCTOR_KWARGS = {"kms_key_name": KEY_NAME}

UPLOAD_KWARGS = {"content-type": "text/fake"}

MAX_WORKERS = 7

DEADLINE = 10

WORKER_TYPE = transfer_manager.THREAD

EXPECTED_FILE_BLOB_PAIRS = [

(os.path.join(ROOT, filename), mock.ANY) for filename in FILENAMES

]

with mock.patch(

"google.cloud.storage.transfer_manager.upload_many"

) as mock_upload_many:

transfer_manager.upload_many_from_filenames(

bucket,

FILENAMES,

source_directory=ROOT,

blob_name_prefix=PREFIX,

skip_if_exists=True,

blob_constructor_kwargs=BLOB_CONSTRUCTOR_KWARGS,

upload_kwargs=UPLOAD_KWARGS,

deadline=DEADLINE,

raise_exception=True,

worker_type=WORKER_TYPE,

max_workers=MAX_WORKERS,

)

mock_upload_many.assert_called_once_with(

EXPECTED_FILE_BLOB_PAIRS,

skip_if_exists=True,

upload_kwargs=UPLOAD_KWARGS,

deadline=DEADLINE,

raise_exception=True,

worker_type=WORKER_TYPE,

max_workers=MAX_WORKERS,

)

bucket.blob.assert_any_call(PREFIX + FILENAMES[0], **BLOB_CONSTRUCTOR_KWARGS)

bucket.blob.assert_any_call(PREFIX + FILENAMES[1], **BLOB_CONSTRUCTOR_KWARGS)

def test_upload_many_from_filenames_minimal_args():

bucket = mock.Mock()

FILENAMES = ["file_a.txt", "file_b.txt"]

EXPECTED_FILE_BLOB_PAIRS = [(filename, mock.ANY) for filename in FILENAMES]

with mock.patch(

"google.cloud.storage.transfer_manager.upload_many"

) as mock_upload_many:

transfer_manager.upload_many_from_filenames(

bucket,

FILENAMES,

)

mock_upload_many.assert_called_once_with(

EXPECTED_FILE_BLOB_PAIRS,

skip_if_exists=False,

upload_kwargs=None,

deadline=None,

raise_exception=False,

worker_type=transfer_manager.PROCESS,

max_workers=8,

)

bucket.blob.assert_any_call(FILENAMES[0])

bucket.blob.assert_any_call(FILENAMES[1])

def test_upload_many_from_filenames_additional_properties():

bucket = mock.Mock()

blob = mock.Mock()

bucket_blob = mock.Mock(return_value=blob)

blob.cache_control = None

bucket.blob = bucket_blob

FILENAME = "file_a.txt"

ADDITIONAL_BLOB_ATTRIBUTES = {"cache_control": "no-cache"}

EXPECTED_FILE_BLOB_PAIRS = [(FILENAME, mock.ANY)]

with mock.patch(

"google.cloud.storage.transfer_manager.upload_many"

) as mock_upload_many:

transfer_manager.upload_many_from_filenames(

bucket, [FILENAME], additional_blob_attributes=ADDITIONAL_BLOB_ATTRIBUTES

)

mock_upload_many.assert_called_once_with(

EXPECTED_FILE_BLOB_PAIRS,

skip_if_exists=False,

upload_kwargs=None,

deadline=None,

raise_exception=False,

worker_type=transfer_manager.PROCESS,

max_workers=8,

)

for attrib, value in ADDITIONAL_BLOB_ATTRIBUTES.items():

assert getattr(blob, attrib) == value

def test__resolve_path_raises_invalid_path_error_on_windows():

from google.cloud.storage.transfer_manager import _resolve_path, InvalidPathError

with mock.patch("os.name", "nt"):

with pytest.raises(InvalidPathError) as exc_info:

_resolve_path("C:\\target", "C:\\target\\file.txt")

assert "cannot be downloaded into" in str(exc_info.value)

# Test that it DOES NOT raise on non-windows

with mock.patch("os.name", "posix"):

# Should not raise

_resolve_path("/target", "C:\\target\\file.txt")

def test_download_many_to_path_raises_invalid_path_error():

bucket = mock.Mock()

BLOBNAMES = ["C:\\target\\file.txt"]

PATH_ROOT = "mypath/"

BLOB_NAME_PREFIX = "myprefix/"

DOWNLOAD_KWARGS = {"accept-encoding": "fake-gzip"}

MAX_WORKERS = 7

DEADLINE = 10

WORKER_TYPE = transfer_manager.THREAD

with mock.patch("os.name", "nt"):

import warnings

with warnings.catch_warnings(record=True) as w:

warnings.simplefilter("always")

results = transfer_manager.download_many_to_path(

bucket,

BLOBNAMES,

destination_directory=PATH_ROOT,

blob_name_prefix=BLOB_NAME_PREFIX,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

create_directories=False,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=True,

)

assert len(w) == 1

assert "will **NOT** be downloaded" in str(w[0].message)

assert len(results) == 1

assert isinstance(results[0], UserWarning)

def test_download_many_to_path():

bucket = mock.Mock()

BLOBNAMES = ["file_a.txt", "file_b.txt", "dir_a/file_c.txt"]

PATH_ROOT = "mypath/"

BLOB_NAME_PREFIX = "myprefix/"

DOWNLOAD_KWARGS = {"accept-encoding": "fake-gzip"}

MAX_WORKERS = 7

DEADLINE = 10

WORKER_TYPE = transfer_manager.THREAD

EXPECTED_BLOB_FILE_PAIRS = [

(mock.ANY, os.path.join(os.getcwd(), PATH_ROOT, blobname))

for blobname in BLOBNAMES

]

with mock.patch(

"google.cloud.storage.transfer_manager.download_many",

return_value=[FAKE_RESULT] * len(BLOBNAMES),

) as mock_download_many:

results = transfer_manager.download_many_to_path(

bucket,

BLOBNAMES,

destination_directory=PATH_ROOT,

blob_name_prefix=BLOB_NAME_PREFIX,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

create_directories=False,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=True,

)

mock_download_many.assert_called_once_with(

EXPECTED_BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=False,

)

assert results == [FAKE_RESULT] * len(BLOBNAMES)

for blobname in BLOBNAMES:

bucket.blob.assert_any_call(BLOB_NAME_PREFIX + blobname)

def test_download_many_to_path_with_skip_if_exists():

bucket = mock.Mock()

BLOBNAMES = ["file_a.txt", "file_b.txt", "dir_a/file_c.txt"]

PATH_ROOT = "mypath/"

BLOB_NAME_PREFIX = "myprefix/"

DOWNLOAD_KWARGS = {"accept-encoding": "fake-gzip"}

MAX_WORKERS = 7

DEADLINE = 10

WORKER_TYPE = transfer_manager.THREAD

from google.cloud.storage.transfer_manager import _resolve_path

existing_file = str(_resolve_path(PATH_ROOT, "file_a.txt"))

def isfile_side_effect(path):

return path == existing_file

EXPECTED_BLOB_FILE_PAIRS = [

(mock.ANY, str(_resolve_path(PATH_ROOT, "file_b.txt"))),

(mock.ANY, str(_resolve_path(PATH_ROOT, "dir_a/file_c.txt"))),

]

with mock.patch("os.path.isfile", side_effect=isfile_side_effect):

with mock.patch(

"google.cloud.storage.transfer_manager.download_many",

return_value=[FAKE_RESULT, FAKE_RESULT],

) as mock_download_many:

results = transfer_manager.download_many_to_path(

bucket,

BLOBNAMES,

destination_directory=PATH_ROOT,

blob_name_prefix=BLOB_NAME_PREFIX,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

create_directories=False,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=True,

)

mock_download_many.assert_called_once_with(

EXPECTED_BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=False,

)

assert len(results) == 3

assert isinstance(results[0], UserWarning)

assert str(results[0]) == "The blob file_a.txt is skipped because destination file already exists"

assert results[1] == FAKE_RESULT

assert results[2] == FAKE_RESULT

@pytest.mark.parametrize(

"blobname",

[

"../../local/target",

"../mypath",

"../escape.txt",

"go/four/levels/deep/../../../../../somefile1",

"go/four/levels/deep/../some_dir/../../../../../invalid/path1",

],

)

def test_download_many_to_path_skips_download(blobname):

bucket = mock.Mock()

BLOBNAMES = [blobname]

PATH_ROOT = "mypath/"

BLOB_NAME_PREFIX = "myprefix/"

DOWNLOAD_KWARGS = {"accept-encoding": "fake-gzip"}

MAX_WORKERS = 7

DEADLINE = 10

WORKER_TYPE = transfer_manager.THREAD

import warnings

with warnings.catch_warnings(record=True) as w:

warnings.simplefilter("always")

with mock.patch(

"google.cloud.storage.transfer_manager.download_many",

return_value=[],

) as mock_download_many:

results = transfer_manager.download_many_to_path(

bucket,

BLOBNAMES,

destination_directory=PATH_ROOT,

blob_name_prefix=BLOB_NAME_PREFIX,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

create_directories=False,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=True,

)

path_traversal_warnings = [

warning for warning in w

if str(warning.message).startswith("The blob ")

and "will **NOT** be downloaded. The resolved destination_directory" in str(warning.message)

]

assert len(path_traversal_warnings) == 1, "---".join([str(warning.message) for warning in w])

mock_download_many.assert_called_once_with(

[],

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=False,

)

assert len(results) == 1

assert isinstance(results[0], UserWarning)

@pytest.mark.parametrize(

"blobname",

[

"simple_blob",

"data/file.txt",

"data/../sibling.txt",

"/etc/passwd",

"/local/usr/a.txt",

"dir/./file.txt",

"go/four/levels/deep/../somefile2",

"go/four/levels/deep/../some_dir/valid/path1",

"go/four/levels/deep/../some_dir/../../../../valid/path2",

],

)

def test_download_many_to_path_downloads_within_dest_dir(blobname):

bucket = mock.Mock()

BLOBNAMES = [blobname]

PATH_ROOT = "mypath/"

BLOB_NAME_PREFIX = "myprefix/"

DOWNLOAD_KWARGS = {"accept-encoding": "fake-gzip"}

MAX_WORKERS = 7

DEADLINE = 10

WORKER_TYPE = transfer_manager.THREAD

from google.cloud.storage.transfer_manager import _resolve_path

EXPECTED_BLOB_FILE_PAIRS = [

(mock.ANY, str(_resolve_path(PATH_ROOT, blobname)))

]

with mock.patch(

"google.cloud.storage.transfer_manager.download_many",

return_value=[FAKE_RESULT],

) as mock_download_many:

results = transfer_manager.download_many_to_path(

bucket,

BLOBNAMES,

destination_directory=PATH_ROOT,

blob_name_prefix=BLOB_NAME_PREFIX,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

create_directories=False,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=True,

)

mock_download_many.assert_called_once_with(

EXPECTED_BLOB_FILE_PAIRS,

download_kwargs=DOWNLOAD_KWARGS,

deadline=DEADLINE,

raise_exception=True,

max_workers=MAX_WORKERS,

worker_type=WORKER_TYPE,

skip_if_exists=False,

)

assert results == [FAKE_RESULT]

bucket.blob.assert_any_call(BLOB_NAME_PREFIX + blobname)

def test_download_many_to_path_creates_directories():

bucket = mock.Mock()

with tempfile.TemporaryDirectory() as tempdir:

DIR_NAME = "dir_a/dir_b"

BLOBNAMES = [

"file_a.txt",

"file_b.txt",

os.path.join(DIR_NAME, "file_c.txt"),

]

EXPECTED_BLOB_FILE_PAIRS = [

(mock.ANY, os.path.join(tempdir, blobname)) for blobname in BLOBNAMES

]

with mock.patch(

"google.cloud.storage.transfer_manager.download_many"

) as mock_download_many:

transfer_manager.download_many_to_path(

bucket,

BLOBNAMES,

destination_directory=tempdir,

create_directories=True,

raise_exception=True,

)

mock_download_many.assert_called_once_with(

EXPECTED_BLOB_FILE_PAIRS,

download_kwargs=None,

deadline=None,

raise_exception=True,

worker_type=transfer_manager.PROCESS,

max_workers=8,

skip_if_exists=False,

)

for blobname in BLOBNAMES:

bucket.blob.assert_any_call(blobname)

assert os.path.isdir(os.path.join(tempdir, DIR_NAME))

def test_download_chunks_concurrently():

blob_mock = mock.Mock(spec=Blob)

FILENAME = "file_a.txt"

MULTIPLE = 4

blob_mock.size = CHUNK_SIZE * MULTIPLE

expected_download_kwargs = EXPECTED_DOWNLOAD_KWARGS.copy()

expected_download_kwargs["command"] = "tm.download_sharded"

expected_download_kwargs["checksum"] = None

with mock.patch("google.cloud.storage.transfer_manager.open", mock.mock_open()):

result = transfer_manager.download_chunks_concurrently(

blob_mock,

FILENAME,

chunk_size=CHUNK_SIZE,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

crc32c_checksum=False,

)

for x in range(MULTIPLE):

blob_mock._prep_and_do_download.assert_any_call(

mock.ANY,

**expected_download_kwargs,

start=x * CHUNK_SIZE,

end=((x + 1) * CHUNK_SIZE) - 1,

)

assert blob_mock._prep_and_do_download.call_count == 4

assert result is None

def test_download_chunks_concurrently_with_crc32c():

blob_mock = mock.Mock(spec=Blob)

FILENAME = "file_a.txt"

MULTIPLE = 4

BLOB_CHUNK = b"abcdefgh"

BLOB_CONTENTS = BLOB_CHUNK * MULTIPLE

blob_mock.size = len(BLOB_CONTENTS)

blob_mock.crc32c = "eOVVVw=="

def write_to_file(f, *args, **kwargs):

f.write(BLOB_CHUNK)

blob_mock._prep_and_do_download.side_effect = write_to_file

with mock.patch("google.cloud.storage.transfer_manager.open", mock.mock_open()):

transfer_manager.download_chunks_concurrently(

blob_mock,

FILENAME,

chunk_size=CHUNK_SIZE,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

crc32c_checksum=True,

)

def test_download_chunks_concurrently_with_crc32c_failure():

blob_mock = mock.Mock(spec=Blob)

FILENAME = "file_a.txt"

MULTIPLE = 4

BLOB_CHUNK = b"abcdefgh"

BLOB_CONTENTS = BLOB_CHUNK * MULTIPLE

blob_mock.size = len(BLOB_CONTENTS)

blob_mock.crc32c = "invalid"

def write_to_file(f, *args, **kwargs):

f.write(BLOB_CHUNK)

blob_mock._prep_and_do_download.side_effect = write_to_file

with mock.patch("google.cloud.storage.transfer_manager.open", mock.mock_open()):

with pytest.raises(DataCorruption):

transfer_manager.download_chunks_concurrently(

blob_mock,

FILENAME,

chunk_size=CHUNK_SIZE,

download_kwargs=DOWNLOAD_KWARGS,

worker_type=transfer_manager.THREAD,

crc32c_checksum=True,

)

def test_download_chunks_concurrently_raises_on_invalid_kwargs():

blob_mock = mock.Mock(spec=Blob)

FILENAME = "file_a.txt"

MULTIPLE = 4

blob_mock.size = CHUNK_SIZE * MULTIPLE

with mock.patch("google.cloud.storage.transfer_manager.open", mock.mock_open()):

with pytest.raises(ValueError):

transfer_manager.download_chunks_concurrently(

blob_mock,

FILENAME,

chunk_size=CHUNK_SIZE,

worker_type=transfer_manager.THREAD,

download_kwargs={

"start": CHUNK_SIZE,

},

)

with pytest.raises(ValueError):

transfer_manager.download_chunks_concurrently(

blob_mock,

FILENAME,

chunk_size=CHUNK_SIZE,

worker_type=transfer_manager.THREAD,

download_kwargs={

"end": (CHUNK_SIZE * (MULTIPLE - 1)) - 1,

},

)

with pytest.raises(ValueError):

transfer_manager.download_chunks_concurrently(

blob_mock,

FILENAME,

chunk_size=CHUNK_SIZE,

worker_type=transfer_manager.THREAD,

download_kwargs={

"checksum": "crc32c",

},

)

def test_download_chunks_concurrently_passes_concurrency_options():

blob_mock = mock.Mock(spec=Blob)

FILENAME = "file_a.txt"

MAX_WORKERS = 7

DEADLINE = 10

MULTIPLE = 4

blob_mock.size = CHUNK_SIZE * MULTIPLE

with mock.patch("concurrent.futures.ThreadPoolExecutor") as pool_patch, mock.patch(

"concurrent.futures.wait"

) as wait_patch, mock.patch(

"google.cloud.storage.transfer_manager.open", mock.mock_open()

):

transfer_manager.download_chunks_concurrently(

blob_mock,

FILENAME,

chunk_size=CHUNK_SIZE,

deadline=DEADLINE,

worker_type=transfer_manager.THREAD,

max_workers=MAX_WORKERS,

crc32c_checksum=False,

)

pool_patch.assert_called_with(max_workers=MAX_WORKERS)

wait_patch.assert_called_with(mock.ANY, timeout=DEADLINE, return_when=mock.ANY)

def test_upload_chunks_concurrently():

bucket = mock.Mock()

bucket.name = "bucket"

bucket.client = _PickleableMockClient(identify_as_client=True)

transport = bucket.client._http

bucket.user_project = None

blob = Blob("blob", bucket)

blob.content_type = FAKE_CONTENT_TYPE

FILENAME = "file_a.txt"

SIZE = 2048

container_mock = mock.Mock()

container_mock.upload_id = "abcd"

part_mock = mock.Mock()

ETAG = "efgh"