From 7c0c3930a8f2d38bb44d22b7633cfe8738d817c6 Mon Sep 17 00:00:00 2001 From: Rhea Parekh Date: Sun, 15 Apr 2018 19:25:41 +0530 Subject: [PATCH] slack importer: Thread avatar downloads. --- zerver/lib/slack_data_to_zulip_data.py | 31 ++++++++++++++++---------- zerver/tests/test_slack_importer.py | 22 ------------------ 2 files changed, 19 insertions(+), 34 deletions(-) diff --git a/zerver/lib/slack_data_to_zulip_data.py b/zerver/lib/slack_data_to_zulip_data.py index fa84ae2dac..1c93e4a234 100755 --- a/zerver/lib/slack_data_to_zulip_data.py +++ b/zerver/lib/slack_data_to_zulip_data.py @@ -956,9 +956,21 @@ def process_avatars(avatar_list: List[ZerverFieldsT], avatar_dir: str, user's avatar directory with both the extensions '.png' and '.original' """ + def get_avatar(avatar_upload_list: List[str]) -> int: + # get avatar of size 512 + slack_avatar_url = avatar_upload_list[0] + image_path = avatar_upload_list[1] + original_image_path = avatar_upload_list[2] + response = requests.get(slack_avatar_url + '-512', stream=True) + with open(image_path, 'wb') as image_file: + shutil.copyfileobj(response.raw, image_file) + shutil.copy(image_path, original_image_path) + return 0 + logging.info('######### GETTING AVATARS #########\n') logging.info('DOWNLOADING AVATARS .......\n') avatar_original_list = [] + avatar_upload_list = [] for avatar in avatar_list: avatar_hash = user_avatar_path_from_ids(avatar['user_profile_id'], realm_id) slack_avatar_url = avatar['path'] @@ -967,28 +979,23 @@ def process_avatars(avatar_list: List[ZerverFieldsT], avatar_dir: str, image_path = ('%s/%s.png' % (avatar_dir, avatar_hash)) original_image_path = ('%s/%s.original' % (avatar_dir, avatar_hash)) - # Fetch the avatars from the url - get_avatar(slack_avatar_url, image_path, original_image_path) - image_size = os.stat(image_path).st_size + avatar_upload_list.append([slack_avatar_url, image_path, original_image_path]) avatar['path'] = image_path avatar['s3_path'] = image_path - avatar['size'] = image_size avatar_original['path'] = original_image_path avatar_original['s3_path'] = original_image_path - avatar_original['size'] = image_size avatar_original_list.append(avatar_original) + + # Run downloads parallely + output = [] + for (status, job) in run_parallel(get_avatar, avatar_upload_list): + output.append(job) + logging.info('######### GETTING AVATARS FINISHED #########\n') return avatar_list + avatar_original_list -def get_avatar(slack_avatar_url: str, image_path: str, original_image_path: str) -> None: - # get avatar of size 512 - response = requests.get(slack_avatar_url + '-512', stream=True) - with open(image_path, 'wb') as image_file: - shutil.copyfileobj(response.raw, image_file) - shutil.copy(image_path, original_image_path) - def process_uploads(upload_list: List[ZerverFieldsT], upload_dir: str) -> List[ZerverFieldsT]: """ This function gets the uploads and saves it in the realm's upload directory diff --git a/zerver/tests/test_slack_importer.py b/zerver/tests/test_slack_importer.py index 310732c0e5..a6f227e37f 100644 --- a/zerver/tests/test_slack_importer.py +++ b/zerver/tests/test_slack_importer.py @@ -104,28 +104,6 @@ class SlackImporter(ZulipTestCase): self.assertEqual(test_zerver_realm_dict['name'], realm_subdomain) self.assertEqual(test_zerver_realm_dict['date_created'], time) - @mock.patch("zerver.lib.slack_data_to_zulip_data.get_avatar") - @mock.patch("os.stat") - def test_user_avatars(self, mock_get_avatar: mock.Mock, mock_os: mock.Mock) -> None: - avatar_url = "https://ca.slack-edge.com/{}-{}-{}".format('T5YFFM2QY', 'U6006P1CN', - 'gd41c3c33cbe') - self.assertEqual(build_avatar_url('U6006P1CN', 'T5YFFM2QY', 'gd41c3c33cbe'), avatar_url) - - avatar_list = [] # type: List[Dict[str, Any]] - timestamp = int(timezone_now().timestamp()) - build_avatar(1, 1, 'email', avatar_url, timestamp, avatar_list) - self.assertEqual(avatar_list[0]['path'], avatar_url) - self.assertEqual(avatar_list[0]['s3_path'], '') - self.assertEqual(avatar_list[0]['user_profile_id'], 1) - - avatar_list = process_avatars(avatar_list, './avatar_dir', 3) - avatar_hash = user_avatar_path_from_ids(1, 3) - image_path = ('%s/%s.png' % ('./avatar_dir', avatar_hash)) - original_image_path = ('%s/%s.original' % ('./avatar_dir', avatar_hash)) - - self.assertEqual(avatar_list[0]['path'], image_path) - self.assertEqual(avatar_list[1]['path'], original_image_path) - def test_get_admin(self) -> None: user_data = [{'is_admin': True, 'is_owner': False, 'is_primary_owner': False}, {'is_admin': True, 'is_owner': True, 'is_primary_owner': False},