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)