def test_copy_file_from_worker(self):
"""Test file_host.copy_file_from_worker."""
mock_response = mock.MagicMock()
mock_response.trailing_metadata.return_value = (('result', 'ok'), )
mock_response.__iter__.return_value = iter([
untrusted_runner_pb2.FileChunk(data='A'),
untrusted_runner_pb2.FileChunk(data='B'),
untrusted_runner_pb2.FileChunk(data='C'),
])
self.mock.stub().CopyFileFrom.return_value = mock_response
self.assertTrue(file_host.copy_file_from_worker('/file', '/file'))
with open('/file') as f:
self.assertEqual(f.read(), 'ABC')
def test_copy_file_to_worker(self):
"""Test file_impl.copy_file_to_worker."""
request_iterator = (
untrusted_runner_pb2.FileChunk(data='A'),
untrusted_runner_pb2.FileChunk(data='B'),
untrusted_runner_pb2.FileChunk(data='C'),
)
context = mock.MagicMock()
context.invocation_metadata.return_value = (('path-bin', '/file'), )
response = file_impl.copy_file_to_worker(request_iterator, context)
self.assertTrue(response.result)
self.assertTrue(os.path.exists('/file'))
with open('/file') as f:
self.assertEqual('ABC', f.read())
def test_copy_file_to_worker_create_dir_is_a_file(self):
"""Test file_impl.copy_file_to_worker when the directory is an existing
file."""
request_iterator = (
untrusted_runner_pb2.FileChunk(data=b'A'),
untrusted_runner_pb2.FileChunk(data=b'B'),
untrusted_runner_pb2.FileChunk(data=b'C'),
)
self.fs.create_file('/file')
context = mock.MagicMock()
context.invocation_metadata.return_value = (('path-bin', b'/file/file'),)
response = file_impl.copy_file_to_worker(request_iterator, context)
self.assertFalse(response.result)
self.assertTrue(os.path.isfile('/file'))
def data_chunk_generator(data):
"""Yields chunks for data."""
index = 0
while index < len(data):
cur_chunk = data[index:index + config.FILE_TRANSFER_CHUNK_SIZE]
yield untrusted_runner_pb2.FileChunk(data=cur_chunk)
index += config.FILE_TRANSFER_CHUNK_SIZE
def test_copy_file_to_worker_create_dir_error(self):
"""Test file_impl.copy_file_to_worker when we fail to create intermediate
dirs."""
request_iterator = (
untrusted_runner_pb2.FileChunk(data='A'),
untrusted_runner_pb2.FileChunk(data='B'),
untrusted_runner_pb2.FileChunk(data='C'),
)
self.fs.CreateFile('/file')
context = mock.MagicMock()
context.invocation_metadata.return_value = (('path-bin',
'/file/dir/file'), )
response = file_impl.copy_file_to_worker(request_iterator, context)
self.assertFalse(response.result)
self.assertTrue(os.path.isfile('/file'))
def file_chunk_generator(handle):
"""Yields chunks from handle."""
data = handle.read(config.FILE_TRANSFER_CHUNK_SIZE)
while data:
yield untrusted_runner_pb2.FileChunk(data=data)
data = handle.read(config.FILE_TRANSFER_CHUNK_SIZE)
