upload: Create function to delete avatar image.

zerver/lib/actions.py

@@ -126,7 +126,7 @@ from zerver.lib.narrow import check_supported_events_narrow_filter
 from zerver.lib.exceptions import JsonableError, ErrorCode, BugdownRenderingException
 from zerver.lib.sessions import delete_user_sessions
 from zerver.lib.upload import attachment_url_re, attachment_url_to_path_id, \
-    claim_attachment, delete_message_image, upload_emoji_image
+    claim_attachment, delete_message_image, upload_emoji_image, delete_avatar_image
 from zerver.lib.str_utils import NonBinaryStr
 from zerver.tornado.event_queue import request_event_queue, send_event
 from zerver.lib.types import ProfileFieldData
@@ -3021,6 +3021,9 @@ def do_change_avatar_fields(user_profile: UserProfile, avatar_source: str) -> No
                     person=payload),
                active_user_ids(user_profile.realm_id))
 
+def do_delete_avatar_image(user: UserProfile) -> None:
+    do_change_avatar_fields(user, UserProfile.AVATAR_FROM_GRAVATAR)
+    delete_avatar_image(user)
 
 def do_change_icon_source(realm: Realm, icon_source: str, log: bool=True) -> None:
     realm.icon_source = icon_source

zerver/lib/upload.py

@@ -167,6 +167,9 @@ class ZulipUploadBackend:
                             target_user_profile: UserProfile) -> None:
         raise NotImplementedError()
 
+    def delete_avatar_image(self, user: UserProfile) -> None:
+        raise NotImplementedError()
+
     def delete_message_image(self, path_id: str) -> bool:
         raise NotImplementedError()
 
@@ -359,6 +362,14 @@ class S3UploadBackend(ZulipUploadBackend):
         self.write_avatar_images(s3_file_name, target_user_profile,
                                  image_data, content_type)
 
+    def delete_avatar_image(self, user: UserProfile) -> None:
+        path_id = user_avatar_path(user)
+        bucket_name = settings.S3_AVATAR_BUCKET
+
+        self.delete_file_from_s3(path_id + ".original", bucket_name)
+        self.delete_file_from_s3(path_id + "-medium.png", bucket_name)
+        self.delete_file_from_s3(path_id, bucket_name)
+
     def get_avatar_key(self, file_name: str) -> Key:
         conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
         bucket_name = settings.S3_AVATAR_BUCKET
@@ -531,6 +542,13 @@ class LocalUploadBackend(ZulipUploadBackend):
         image_data = user_file.read()
         self.write_avatar_images(file_path, image_data)
 
+    def delete_avatar_image(self, user: UserProfile) -> None:
+        path_id = user_avatar_path(user)
+
+        delete_local_file("avatars", path_id + ".original")
+        delete_local_file("avatars", path_id + ".png")
+        delete_local_file("avatars", path_id + "-medium.png")
+
     def get_avatar_url(self, hash_key: str, medium: bool=False) -> str:
         # ?x=x allows templates to append additional parameters with &s
         medium_suffix = "-medium" if medium else ""
@@ -607,6 +625,9 @@ def upload_avatar_image(user_file: File, acting_user_profile: UserProfile,
                         target_user_profile: UserProfile) -> None:
     upload_backend.upload_avatar_image(user_file, acting_user_profile, target_user_profile)
 
+def delete_avatar_image(user_profile: UserProfile):
+    upload_backend.delete_avatar_image(user_profile)
+
 def copy_avatar(source_profile: UserProfile, target_profile: UserProfile) -> None:
     upload_backend.copy_avatar(source_profile, target_profile)
 

zerver/tests/test_upload.py

@@ -970,6 +970,29 @@ class AvatarTest(UploadSerializeMixin, ZulipTestCase):
         target_medium_path_id = avatar_disk_path(target_user_profile, medium=True)
         self.assertEqual(open(source_medium_path_id, "rb").read(), open(target_medium_path_id, "rb").read())
 
+    def test_delete_avatar_image(self) -> None:
+        self.login(self.example_email("hamlet"))
+        with get_test_image_file('img.png') as image_file:
+            self.client_post("/json/users/me/avatar", {'file': image_file})
+
+        user = self.example_user('hamlet')
+
+        avatar_path_id = avatar_disk_path(user)
+        avatar_original_path_id = avatar_disk_path(user, original=True)
+        avatar_medium_path_id = avatar_disk_path(user, medium=True)
+
+        self.assertEqual(user.avatar_source, UserProfile.AVATAR_FROM_USER)
+        self.assertTrue(os.path.isfile(avatar_path_id))
+        self.assertTrue(os.path.isfile(avatar_original_path_id))
+        self.assertTrue(os.path.isfile(avatar_medium_path_id))
+
+        zerver.lib.actions.do_delete_avatar_image(user)
+
+        self.assertEqual(user.avatar_source, UserProfile.AVATAR_FROM_GRAVATAR)
+        self.assertFalse(os.path.isfile(avatar_path_id))
+        self.assertFalse(os.path.isfile(avatar_original_path_id))
+        self.assertFalse(os.path.isfile(avatar_medium_path_id))
+
     def test_invalid_avatars(self) -> None:
         """
         A PUT request to /json/users/me/avatar with an invalid file should fail.
@@ -1392,6 +1415,33 @@ class S3Test(ZulipTestCase):
         target_medium_image_data = target_medium_image_key.get_contents_as_string()
         self.assertEqual(source_medium_image_data, target_medium_image_data)
 
+    @use_s3_backend
+    def test_delete_avatar_image(self) -> None:
+        conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
+        bucket = conn.create_bucket(settings.S3_AVATAR_BUCKET)
+
+        self.login(self.example_email("hamlet"))
+        with get_test_image_file('img.png') as image_file:
+            self.client_post("/json/users/me/avatar", {'file': image_file})
+
+        user = self.example_user('hamlet')
+
+        avatar_path_id = user_avatar_path(user)
+        avatar_original_image_path_id = avatar_path_id + ".original"
+        avatar_medium_path_id = avatar_path_id + "-medium.png"
+
+        self.assertEqual(user.avatar_source, UserProfile.AVATAR_FROM_USER)
+        self.assertIsNotNone(bucket.get_key(avatar_path_id))
+        self.assertIsNotNone(bucket.get_key(avatar_original_image_path_id))
+        self.assertIsNotNone(bucket.get_key(avatar_medium_path_id))
+
+        zerver.lib.actions.do_delete_avatar_image(user)
+
+        self.assertEqual(user.avatar_source, UserProfile.AVATAR_FROM_GRAVATAR)
+        self.assertIsNone(bucket.get_key(avatar_path_id))
+        self.assertIsNone(bucket.get_key(avatar_original_image_path_id))
+        self.assertIsNone(bucket.get_key(avatar_medium_path_id))
+
     @use_s3_backend
     def test_get_realm_for_filename(self) -> None:
         conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)