def test_copy_file_from_worker_failed(self):
"""Test file_impl.copy_file_from_worker."""
request = untrusted_runner_pb2.CopyFileFromRequest(path='/file')
context = mock.MagicMock()
response = file_impl.copy_file_from_worker(request, context)
self.assertEqual(0, len(list(response)))
context.set_trailing_metadata.assert_called_with([('result',
'invalid-path')])
def copy_file_from_worker(worker_path, host_path):
"""Copy file from worker to host."""
request = untrusted_runner_pb2.CopyFileFromRequest(path=worker_path)
response = host.stub().CopyFileFrom(request)
file_utils.write_chunks(host_path, response)
metadata = dict(response.trailing_metadata())
if metadata.get('result') != 'ok':
# file_utils.write_chunks always opens the file for writing, so remove it
# here.
os.remove(host_path)
return False
return True
def test_copy_file_from_worker(self):
"""Test file_impl.copy_file_from_worker."""
contents = ('A' * config.FILE_TRANSFER_CHUNK_SIZE +
'B' * config.FILE_TRANSFER_CHUNK_SIZE +
'C' * config.FILE_TRANSFER_CHUNK_SIZE)
self.fs.CreateFile('/file', contents=contents)
request = untrusted_runner_pb2.CopyFileFromRequest(path='/file')
context = mock.MagicMock()
response = file_impl.copy_file_from_worker(request, context)
chunks = [chunk.data for chunk in response]
self.assertEqual(len(chunks), 3)
self.assertEqual(contents, ''.join(chunks))
context.set_trailing_metadata.assert_called_with([('result', 'ok')])