def test_pull_from_stream(self):
s3handler = S3StreamHandler(self.session, self.params, chunksize=2)
input_to_stdin = b'This is a test'
size = len(input_to_stdin)
# Retrieve the entire string.
with MockStdIn(input_to_stdin):
payload, is_amount_requested = s3handler._pull_from_stream(size)
data = payload.read()
self.assertTrue(is_amount_requested)
self.assertEqual(data, input_to_stdin)
# Ensure the function exits when there is nothing to read.
with MockStdIn():
payload, is_amount_requested = s3handler._pull_from_stream(size)
data = payload.read()
self.assertFalse(is_amount_requested)
self.assertEqual(data, b'')
# Ensure the function does not grab too much out of stdin.
with MockStdIn(input_to_stdin):
payload, is_amount_requested = s3handler._pull_from_stream(size-2)
data = payload.read()
self.assertTrue(is_amount_requested)
self.assertEqual(data, input_to_stdin[:-2])
# Retrieve the rest of standard in.
payload, is_amount_requested = s3handler._pull_from_stream(size)
data = payload.read()
self.assertFalse(is_amount_requested)
self.assertEqual(data, input_to_stdin[-2:])
def test_enqueue_multipart_download_stream(self):
"""
This test ensures the right calls are made in ``_enqueue_tasks()``
if the file should be a multipart download.
"""
s3handler = S3StreamHandler(self.session, self.params,
multi_threshold=5)
s3handler.executor = mock.Mock()
fileinfo = FileInfo('filename', operation_name='download',
is_stream=True)
with mock.patch('awscli.customizations.s3.s3handler'
'.S3StreamHandler._enqueue_range_download_tasks') as \
mock_enqueue_range_tasks:
with mock.patch('awscli.customizations.s3.fileinfo.FileInfo'
'.set_size_from_s3') as mock_set_size_from_s3:
# Set the file size to something larger than the multipart
# threshold.
fileinfo.size = 100
# Run the main enqueue function.
s3handler._enqueue_tasks([fileinfo])
# Assert that the size of the ``FileInfo`` object was set
# if we are downloading a stream.
self.assertTrue(mock_set_size_from_s3.called)
# Ensure that this download would have been a multipart
# download.
self.assertTrue(mock_enqueue_range_tasks.called)
def test_enqueue_range_download_tasks_stream(self):
s3handler = S3StreamHandler(self.session, self.params, chunksize=100)
s3handler.executor = mock.Mock()
fileinfo = FileInfo('filename', operation_name='download',
is_stream=True, size=100)
s3handler._enqueue_range_download_tasks(fileinfo)
# Ensure that no request was sent to make a file locally.
submitted_tasks = s3handler.executor.submit.call_args_list
self.assertNotEqual(type(submitted_tasks[0][0][0]),
CreateLocalFileTask)
def test_upload_stream_not_multipart_task(self):
s3handler = S3StreamHandler(self.session, self.params)
s3handler.executor = mock.Mock()
fileinfos = [FileInfo('filename', operation_name='upload',
is_stream=True, size=0)]
with MockStdIn(b'bar'):
s3handler._enqueue_tasks(fileinfos)
submitted_tasks = s3handler.executor.submit.call_args_list
# No multipart upload should have been submitted.
self.assertEqual(len(submitted_tasks), 1)
self.assertEqual(submitted_tasks[0][0][0].payload.read(),
b'bar')
def test_enqueue_upload_single_part_task_stream(self):
"""
This test ensures that a payload gets attached to a task when
it is submitted to the executor.
"""
s3handler = S3StreamHandler(self.session, self.params)
s3handler.executor = mock.Mock()
mock_task_class = mock.Mock()
s3handler._enqueue_upload_single_part_task(
part_number=1, chunk_size=2, upload_context=None,
filename=None, task_class=mock_task_class,
payload=b'This is a test'
)
args, kwargs = mock_task_class.call_args
self.assertIn('payload', kwargs.keys())
self.assertEqual(kwargs['payload'], b'This is a test')
def test_upload_stream_with_expected_size(self):
self.params['expected_size'] = 100000
# With this large of expected size, the chunksize of 2 will have
# to change.
s3handler = S3StreamHandler(self.session, self.params, chunksize=2)
s3handler.executor = mock.Mock()
fileinfo = FileInfo('filename', operation_name='upload',
is_stream=True)
with MockStdIn(b'bar'):
s3handler._enqueue_multipart_upload_tasks(fileinfo, b'')
submitted_tasks = s3handler.executor.submit.call_args_list
# Determine what the chunksize was changed to from one of the
# UploadPartTasks.
changed_chunk_size = submitted_tasks[1][0][0]._chunk_size
# New chunksize should have a total parts under 1000.
self.assertTrue(100000/changed_chunk_size < 1000)
def test_upload_stream_is_multipart_task(self):
s3handler = S3StreamHandler(
self.session, self.params,
runtime_config=runtime_config(multipart_threshold=1))
s3handler.executor = mock.Mock()
fileinfos = [FileInfo('filename', operation_name='upload',
is_stream=True, size=0)]
with MockStdIn(b'bar'):
s3handler._enqueue_tasks(fileinfos)
submitted_tasks = s3handler.executor.submit.call_args_list
# This should be a multipart upload so multiple tasks
# should have been submitted.
self.assertEqual(len(submitted_tasks), 4)
self.assertEqual(submitted_tasks[1][0][0]._payload.read(),
b'b')
self.assertEqual(submitted_tasks[2][0][0]._payload.read(),
b'ar')
def test_upload_stream_enqueue_upload_task(self):
s3handler = S3StreamHandler(self.session, self.params)
s3handler.executor = mock.Mock()
fileinfo = FileInfo('filename', operation_name='upload',
is_stream=True)
stdin_input = b'This is a test'
with MockStdIn(stdin_input):
num_parts = s3handler._enqueue_upload_tasks(None, 2, mock.Mock(),
fileinfo,
UploadPartTask)
submitted_tasks = s3handler.executor.submit.call_args_list
# Ensure the returned number of parts is correct.
self.assertEqual(num_parts, len(submitted_tasks) + 1)
# Ensure the number of tasks uploaded are as expected
self.assertEqual(len(submitted_tasks), 8)
index = 0
for i in range(len(submitted_tasks)-1):
self.assertEqual(submitted_tasks[i][0][0]._payload.read(),
stdin_input[index:index+2])
index += 2
# Ensure that the last part is an empty string as expected.
self.assertEqual(submitted_tasks[7][0][0]._payload.read(), b'')
def run(self):
"""
This function wires together all of the generators and completes
the command. First a dictionary is created that is indexed first by
the command name. Then using the instruction, another dictionary
can be indexed to obtain the objects corresponding to the
particular instruction for that command. To begin the wiring,
either a ``FileFormat`` or ``TaskInfo`` object, depending on the
command, is put into a list. Then the function enters a while loop
that pops off an instruction. It then determines the object needed
and calls the call function of the object using the list as the input.
Depending on the number of objects in the input list and the number
of components in the list corresponding to the instruction, the call
method of the component can be called two different ways. If the
number of inputs is equal to the number of components a 1:1 mapping of
inputs to components is used when calling the call function. If the
there are more inputs than components, then a 2:1 mapping of inputs to
components is used where the component call method takes two inputs
instead of one. Whatever files are yielded from the call function
is appended to a list and used as the input for the next repetition
of the while loop until there are no more instructions.
"""
src = self.parameters['src']
dest = self.parameters['dest']
paths_type = self.parameters['paths_type']
files = FileFormat().format(src, dest, self.parameters)
rev_files = FileFormat().format(dest, src, self.parameters)
cmd_translation = {}
cmd_translation['locals3'] = {
'cp': 'upload',
'sync': 'upload',
'mv': 'move'
}
cmd_translation['s3s3'] = {'cp': 'copy', 'sync': 'copy', 'mv': 'move'}
cmd_translation['s3local'] = {
'cp': 'download',
'sync': 'download',
'mv': 'move'
}
cmd_translation['s3'] = {
'rm': 'delete',
'mb': 'make_bucket',
'rb': 'remove_bucket'
}
result_queue = queue.Queue()
operation_name = cmd_translation[paths_type][self.cmd]
file_generator = FileGenerator(self._service,
self._source_endpoint,
operation_name,
self.parameters['follow_symlinks'],
self.parameters['page_size'],
result_queue=result_queue)
rev_generator = FileGenerator(self._service,
self._endpoint,
'',
self.parameters['follow_symlinks'],
self.parameters['page_size'],
result_queue=result_queue)
taskinfo = [
TaskInfo(src=files['src']['path'],
src_type='s3',
operation_name=operation_name,
service=self._service,
endpoint=self._endpoint)
]
stream_dest_path, stream_compare_key = find_dest_path_comp_key(files)
stream_file_info = [
FileInfo(src=files['src']['path'],
dest=stream_dest_path,
compare_key=stream_compare_key,
src_type=files['src']['type'],
dest_type=files['dest']['type'],
operation_name=operation_name,
service=self._service,
endpoint=self._endpoint,
is_stream=True)
]
file_info_builder = FileInfoBuilder(self._service, self._endpoint,
self._source_endpoint,
self.parameters)
s3handler = S3Handler(self.session,
self.parameters,
runtime_config=self._runtime_config,
result_queue=result_queue)
s3_stream_handler = S3StreamHandler(self.session,
self.parameters,
result_queue=result_queue)
sync_strategies = self.choose_sync_strategies()
command_dict = {}
if self.cmd == 'sync':
command_dict = {
'setup': [files, rev_files],
'file_generator': [file_generator, rev_generator],
'filters': [
create_filter(self.parameters),
create_filter(self.parameters)
],
'comparator': [Comparator(**sync_strategies)],
'file_info_builder': [file_info_builder],
's3_handler': [s3handler]
}
elif self.cmd == 'cp' and self.parameters['is_stream']:
command_dict = {
'setup': [stream_file_info],
's3_handler': [s3_stream_handler]
}
elif self.cmd == 'cp':
command_dict = {
'setup': [files],
'file_generator': [file_generator],
'filters': [create_filter(self.parameters)],
'file_info_builder': [file_info_builder],
's3_handler': [s3handler]
}
elif self.cmd == 'rm':
command_dict = {
'setup': [files],
'file_generator': [file_generator],
'filters': [create_filter(self.parameters)],
'file_info_builder': [file_info_builder],
's3_handler': [s3handler]
}
elif self.cmd == 'mv':
command_dict = {
'setup': [files],
'file_generator': [file_generator],
'filters': [create_filter(self.parameters)],
'file_info_builder': [file_info_builder],
's3_handler': [s3handler]
}
elif self.cmd == 'mb':
command_dict = {'setup': [taskinfo], 's3_handler': [s3handler]}
elif self.cmd == 'rb':
command_dict = {'setup': [taskinfo], 's3_handler': [s3handler]}
files = command_dict['setup']
while self.instructions:
instruction = self.instructions.pop(0)
file_list = []
components = command_dict[instruction]
for i in range(len(components)):
if len(files) > len(components):
file_list.append(components[i].call(*files))
else:
file_list.append(components[i].call(files[i]))
files = file_list
# This is kinda quirky, but each call through the instructions
# will replaces the files attr with the return value of the
# file_list. The very last call is a single list of
# [s3_handler], and the s3_handler returns the number of
# tasks failed and the number of tasks warned.
# This means that files[0] now contains a namedtuple with
# the number of failed tasks and the number of warned tasks.
# In terms of the RC, we're keeping it simple and saying
# that > 0 failed tasks will give a 1 RC and > 0 warned
# tasks will give a 2 RC. Otherwise a RC of zero is returned.
rc = 0
if files[0].num_tasks_failed > 0:
rc = 1
if files[0].num_tasks_warned > 0:
rc = 2
return rc
