def testAddTree(self):
# Tree 1
# - folder1 - folder4 - file2
# - file4
# - folder2
# - folder3 - folder5
# - file1.txt
# - file3.png
tree1 = StoreTree(0)
tree1.add_folder(1, 'folder1')
tree1.add_folder(2, 'folder2')
tree1.add_folder(3, 'folder3')
tree1.add_folder(4, 'folder4', parent_id=1)
tree1.add_folder(5, 'folder5', parent_id=3)
tree1.add_file(6, 'file1.txt')
tree1.add_file(7, 'file2', parent_id=4)
tree1.add_file(8, 'file3.png')
tree1.add_file(9, 'file4', parent_id=1)
# Tree 2
# - folder6 - folder7 - file5
# - file6.png
tree2 = StoreTree(100)
tree2.add_folder(10, 'folder6')
tree2.add_folder(11, 'folder7', parent_id=10)
tree2.add_file(12, 'file5.txt', parent_id=11)
tree2.add_file(13, 'file6.png')
# New tree
# - folder1 - folder4 - file2
# - file4
#
# - folder2 - tree2 - folder6 - folder7 - file5.txt
# - file6.png
#
# - folder3 - folder5
# - file1.txt
# - file3.png
tree1.add_tree(tree2, 2)
tree1.update_folder_name(100, 'tree2')
root_dict = tree1.find_item_by_id(0)
self.assertEqual(len(root_dict['files']), 2)
self.assertEqual(len(root_dict['folders']), 3)
self.assertEqual(root_dict['id'], 0)
self.assertEqual(root_dict['name'], '')
tree2_dict = tree1.find_item_by_id(100)
self.assertEqual(len(tree2_dict['files']), 1)
self.assertEqual(len(tree2_dict['folders']), 1)
self.assertEqual(tree2_dict['id'], 100)
self.assertEqual(tree2_dict['name'], 'tree2')
folder6_dict = tree1.find_item_by_id(10)
self.assertEqual(len(folder6_dict['files']), 0)
self.assertEqual(len(folder6_dict['folders']), 1)
self.assertEqual(folder6_dict['id'], 10)
self.assertEqual(folder6_dict['name'], 'folder6')
folder7_dict = tree1.find_item_by_id(11)
self.assertEqual(len(folder7_dict['files']), 1)
self.assertEqual(len(folder7_dict['folders']), 0)
self.assertEqual(folder7_dict['id'], 11)
self.assertEqual(folder7_dict['name'], 'folder7')
file5_dict = folder7_dict['files'][0]
self.assertEqual(file5_dict['id'], 12)
self.assertEqual(file5_dict['name'], 'file5.txt')
file6_dict = tree1.find_item_by_id(13)
self.assertEqual(file6_dict['id'], 13)
self.assertEqual(file6_dict['name'], 'file6.png')
def _get_drive_folder_tree(self):
"""
:return: an instance of StoreTree representing the entire folder tree of the drive.
"""
# Get the root id and create store tree
root_id = self._get_root_folder()['id']
result = StoreTree(id=root_id)
# Google returns items randomly, only specifying the parent id.
# We might not have received the parent item yet, so we maintain
# a list of trees, the first being our result, the others are
# "dangling" trees where the root item hasn't been received yet
# but has been mentioned as a parent of an item that HAS been
# received.
tree_list = [result]
response_dict = None
while response_dict is None or 'nextPageToken' in response_dict:
params = {
'q':
'mimeType = \'application/vnd.google-apps.folder\' and trashed = false',
'fields': 'files/id, files/name, files/parents',
'pageSize': 1000
}
if isinstance(response_dict,
dict) and 'nextPageToken' in response_dict:
params['pageToken'] = response_dict['nextPageToken']
r = self._do_request('get',
http_server_utils.join_url_components([
self._api_drive_endpoint_prefix, 'files'
]),
params=params,
error_500_retries=5)
response_dict = r.json()
for new_folder in response_dict['files']:
# First check if the parent exists in one of the trees tree. If not, we'll
# need to create it as the root of a new dangling tree and update later if/when
# it arrives.
parent_tree = None
for tree in tree_list:
if tree.find_item_by_id(
new_folder['parents'][0]) is not None:
parent_tree = tree
break
if parent_tree is None:
parent_tree = StoreTree(id=new_folder['parents'][0])
tree_list.append(parent_tree)
# Now check if this item has already been added as a parent
# (that will mean it is a tree root). If so, move it to
# its parent and update its name.
added = False
for tree_index in range(0, len(tree_list)):
tree = tree_list[tree_index]
if tree.root_id == new_folder['id']:
tree.update_folder_name(new_folder['id'],
new_folder['name'])
parent_tree.add_tree(tree, new_folder['parents'][0])
del tree_list[tree_index]
added = True
break
# New folder doesn't exist, create a new one.
if added is False:
parent_tree.add_folder(new_folder['id'],
new_folder['name'],
new_folder['parents'][0])
return result