upload: Refactor and add tests for ensure_avatar_image in upload.py.

`ensure_basic_avatar_image` and `ensure_medium_avatar_image` are
essentially the same thing, except a size parameter.
So, refactor them into a single function.

This doesn't introduce any functional changes.

zerver/lib/import_realm.py

@@ -686,9 +686,9 @@ def process_avatars(record: Dict[str, Any]) -> None:
                 # same realm ID from a previous iteration).
                 os.remove(medium_file_path)
         try:
-            upload_backend.ensure_medium_avatar_image(user_profile=user_profile)
+            upload_backend.ensure_avatar_image(user_profile=user_profile, is_medium=True)
             if record.get("importer_should_thumbnail"):
-                upload_backend.ensure_basic_avatar_image(user_profile=user_profile)
+                upload_backend.ensure_avatar_image(user_profile=user_profile)
         except BadImageError:
             logging.warning(
                 "Could not thumbnail avatar image for user %s; ignoring",

zerver/lib/upload.py

@@ -247,10 +247,7 @@ class ZulipUploadBackend:
     def copy_avatar(self, source_profile: UserProfile, target_profile: UserProfile) -> None:
         raise NotImplementedError()
 
-    def ensure_medium_avatar_image(self, user_profile: UserProfile) -> None:
-        raise NotImplementedError()
-
-    def ensure_basic_avatar_image(self, user_profile: UserProfile) -> None:
+    def ensure_avatar_image(self, user_profile: UserProfile, is_medium: bool = False) -> None:
         raise NotImplementedError()
 
     def upload_realm_icon_image(self, icon_file: File, user_profile: UserProfile) -> None:
@@ -590,35 +587,23 @@ class S3UploadBackend(ZulipUploadBackend):
             file_name = "night_logo.png"
         return f"{self.avatar_bucket_url}/{realm_id}/realm/{file_name}?version={version}"
 
-    def ensure_medium_avatar_image(self, user_profile: UserProfile) -> None:
+    def ensure_avatar_image(self, user_profile: UserProfile, is_medium: bool = False) -> None:
+        # BUG: The else case should be user_avatar_path(user_profile) + ".png".
+        # See #12852 for details on this bug and how to migrate it.
+        file_extension = "-medium.png" if is_medium else ""
         file_path = user_avatar_path(user_profile)
         s3_file_name = file_path
 
         key = self.avatar_bucket.Object(file_path + ".original")
         image_data = key.get()["Body"].read()
 
-        resized_medium = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
-        upload_image_to_s3(
-            self.avatar_bucket,
-            s3_file_name + "-medium.png",
-            "image/png",
-            user_profile,
-            resized_medium,
-        )
-
-    def ensure_basic_avatar_image(self, user_profile: UserProfile) -> None:  # nocoverage
-        # TODO: Refactor this to share code with ensure_medium_avatar_image
-        file_path = user_avatar_path(user_profile)
-        # Also TODO: Migrate to user_avatar_path(user_profile) + ".png".
-        s3_file_name = file_path
-
-        key = self.avatar_bucket.Object(file_path + ".original")
-        image_data = key.get()["Body"].read()
-
+        if is_medium:
+            resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
+        else:
             resized_avatar = resize_avatar(image_data)
         upload_image_to_s3(
             self.avatar_bucket,
-            s3_file_name,
+            s3_file_name + file_extension,
             "image/png",
             user_profile,
             resized_avatar,
@@ -859,32 +844,24 @@ class LocalUploadBackend(ZulipUploadBackend):
             file_name = "logo.png"
         return f"/user_avatars/{realm_id}/realm/{file_name}?version={version}"
 
-    def ensure_medium_avatar_image(self, user_profile: UserProfile) -> None:
+    def ensure_avatar_image(self, user_profile: UserProfile, is_medium: bool = False) -> None:
+        file_extension = "-medium.png" if is_medium else ".png"
         file_path = user_avatar_path(user_profile)
 
-        output_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", file_path + "-medium.png")
-        if os.path.isfile(output_path):
-            return
-
-        image_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", file_path + ".original")
-        with open(image_path, "rb") as f:
-            image_data = f.read()
-        resized_medium = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
-        write_local_file("avatars", file_path + "-medium.png", resized_medium)
-
-    def ensure_basic_avatar_image(self, user_profile: UserProfile) -> None:  # nocoverage
-        # TODO: Refactor this to share code with ensure_medium_avatar_image
-        file_path = user_avatar_path(user_profile)
-
-        output_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", file_path + ".png")
+        output_path = os.path.join(
+            settings.LOCAL_UPLOADS_DIR, "avatars", file_path + file_extension
+        )
         if os.path.isfile(output_path):
             return
 
         image_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", file_path + ".original")
         with open(image_path, "rb") as f:
             image_data = f.read()
+        if is_medium:
+            resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
+        else:
             resized_avatar = resize_avatar(image_data)
-        write_local_file("avatars", file_path + ".png", resized_avatar)
+        write_local_file("avatars", file_path + file_extension, resized_avatar)
 
     def upload_emoji_image(
         self, emoji_file: File, emoji_file_name: str, user_profile: UserProfile

zerver/migrations/0032_verify_all_medium_avatar_images.py

@@ -27,7 +27,7 @@ def patched_user_avatar_path(user_profile: UserProfile) -> str:
 def verify_medium_avatar_image(apps: StateApps, schema_editor: DatabaseSchemaEditor) -> None:
     user_profile_model = apps.get_model("zerver", "UserProfile")
     for user_profile in user_profile_model.objects.filter(avatar_source="U"):
-        upload_backend.ensure_medium_avatar_image(user_profile)
+        upload_backend.ensure_avatar_image(user_profile, is_medium=True)
 
 
 class Migration(migrations.Migration):

zerver/tests/test_upload.py

@@ -58,6 +58,7 @@ from zerver.lib.upload import (
     upload_emoji_image,
     upload_export_tarball,
     upload_message_file,
+    write_local_file,
 )
 from zerver.lib.users import get_api_key
 from zerver.models import (
@@ -1129,14 +1130,14 @@ class AvatarTest(UploadSerializeMixin, ZulipTestCase):
 
             # Verify that ensure_medium_avatar_url does not overwrite this file if it exists
             with mock.patch("zerver.lib.upload.write_local_file") as mock_write_local_file:
-                zerver.lib.upload.upload_backend.ensure_medium_avatar_image(user_profile)
+                zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
                 self.assertFalse(mock_write_local_file.called)
 
             # Confirm that ensure_medium_avatar_url works to recreate
             # medium size avatars from the original if needed
             os.remove(medium_avatar_disk_path)
             self.assertFalse(os.path.exists(medium_avatar_disk_path))
-            zerver.lib.upload.upload_backend.ensure_medium_avatar_image(user_profile)
+            zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
             self.assertTrue(os.path.exists(medium_avatar_disk_path))
 
             # Verify whether the avatar_version gets incremented with every new upload
@@ -1641,6 +1642,29 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase):
         path_id = re.sub("/user_uploads/", "", result.json()["uri"])
         self.assertTrue(delete_message_image(path_id))
 
+    def test_ensure_avatar_image_local(self) -> None:
+        user_profile = self.example_user("hamlet")
+        file_path = user_avatar_path(user_profile)
+
+        with get_test_image_file("img.png") as image_file:
+            write_local_file("avatars", file_path + ".original", image_file.read())
+
+        image_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", file_path + ".original")
+        with open(image_path, "rb") as f:
+            image_data = f.read()
+
+        resized_avatar = resize_avatar(image_data)
+        zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile)
+        output_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", file_path + ".png")
+        with open(output_path, "rb") as original_file:
+            self.assertEqual(resized_avatar, original_file.read())
+
+        resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
+        zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
+        output_path = os.path.join(settings.LOCAL_UPLOADS_DIR, "avatars", file_path + "-medium.png")
+        with open(output_path, "rb") as original_file:
+            self.assertEqual(resized_avatar, original_file.read())
+
     def test_emoji_upload_local(self) -> None:
         user_profile = self.example_user("hamlet")
         file_name = "emoji.png"
@@ -1842,7 +1866,7 @@ class S3Test(ZulipTestCase):
         self.assertEqual(medium_image_data, test_medium_image_data)
 
         bucket.Object(medium_image_key.key).delete()
-        zerver.lib.upload.upload_backend.ensure_medium_avatar_image(user_profile)
+        zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
         medium_image_key = bucket.Object(medium_path_id)
         self.assertEqual(medium_image_key.key, medium_path_id)
 
@@ -1991,6 +2015,35 @@ class S3Test(ZulipTestCase):
         resized_image = Image.open(io.BytesIO(resized_data))
         self.assertEqual(resized_image.size, (DEFAULT_EMOJI_SIZE, DEFAULT_EMOJI_SIZE))
 
+    @use_s3_backend
+    def test_ensure_avatar_image(self) -> None:
+        bucket = create_s3_buckets(settings.S3_AVATAR_BUCKET)[0]
+
+        user_profile = self.example_user("hamlet")
+        base_file_path = user_avatar_path(user_profile)
+        # Bug: This should have + ".png", but the implementation is wrong.
+        file_path = base_file_path
+        original_file_path = base_file_path + ".original"
+        medium_file_path = base_file_path + "-medium.png"
+
+        with get_test_image_file("img.png") as image_file:
+            zerver.lib.upload.upload_backend.upload_avatar_image(
+                image_file, user_profile, user_profile
+            )
+
+        key = bucket.Object(original_file_path)
+        image_data = key.get()["Body"].read()
+
+        zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile)
+        resized_avatar = resize_avatar(image_data)
+        key = bucket.Object(file_path)
+        self.assertEqual(resized_avatar, key.get()["Body"].read())
+
+        zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
+        resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
+        key = bucket.Object(medium_file_path)
+        self.assertEqual(resized_avatar, key.get()["Body"].read())
+
     @use_s3_backend
     def test_get_emoji_url(self) -> None:
         emoji_name = "emoji.png"