def copy_file_to_worker(request_iterator, context):
"""Copy file from host to worker."""
metadata = dict(context.invocation_metadata())
path = metadata['path-bin'].decode('utf-8')
# Create intermediate directories if needed.
directory = os.path.dirname(path)
if not os.path.exists(directory):
try:
os.makedirs(directory)
except Exception:
pass
if not os.path.isdir(directory):
# Failed to create intermediate directories.
return untrusted_runner_pb2.CopyFileToResponse(result=False)
file_utils.write_chunks(path, request_iterator)
return untrusted_runner_pb2.CopyFileToResponse(result=True)
def mock_copy_file_to(iterator, metadata):
"""Mock copy file to."""
chunks = [chunk.data for chunk in iterator]
self.assertEqual(3, len(chunks))
self.assertEqual([('path-bin', '/file')], metadata)
data = ''.join(chunks)
self.assertEqual(data, contents)
return untrusted_runner_pb2.CopyFileToResponse(result=True)
def test_write_data_to_worker(self):
"""Test file_host.write_data_to_worker."""
contents = ('A' * config.FILE_TRANSFER_CHUNK_SIZE +
'B' * config.FILE_TRANSFER_CHUNK_SIZE +
'C' * config.FILE_TRANSFER_CHUNK_SIZE)
result = untrusted_runner_pb2.CopyFileToResponse(result=True)
self.mock.stub().CopyFileTo.return_value = result
self.assertTrue(file_host.write_data_to_worker(contents, '/file'))
call_args = self.mock.stub().CopyFileTo.call_args
self.assertEqual(call_args[1], {'metadata': [('path-bin', '/file')]})
chunks = [chunk.data for chunk in call_args[0][0]]
self.assertEqual(len(chunks), 3)
data = ''.join(chunks)
self.assertEqual(data, contents)