chore(hybrid-cloud): Update user avatar tests for split db (#54658)

src/sentry/models/avatars/base.py

@@ -12,6 +12,7 @@ from typing_extensions import Self
 
 from sentry.db.models import BoundedBigIntegerField, Model
 from sentry.models.files.file import File
+from sentry.silo import SiloMode
 from sentry.tasks.files import copy_file_to_control_and_update_model
 from sentry.utils.cache import cache
 from sentry.utils.db import atomic_transaction
@@ -54,14 +55,21 @@ class AvatarBase(Model):
             file_class = File
             if file_id is None:
                 return None
-            copy_file_to_control_and_update_model.apply_async(
-                kwargs={
-                    "app_name": "sentry",
-                    "model_name": type(self).__name__,
-                    "model_id": self.id,
-                    "file_id": file_id,
-                }
-            )
+            if SiloMode.get_current_mode() == SiloMode.MONOLITH:
+                copy_file_to_control_and_update_model.apply_async(
+                    kwargs={
+                        "app_name": "sentry",
+                        "model_name": type(self).__name__,
+                        "model_id": self.id,
+                        "file_id": file_id,
+                    }
+                )
+
+        if (
+            SiloMode.get_current_mode() != SiloMode.MONOLITH
+            and SiloMode.get_current_mode() not in file_class._meta.silo_limit.modes
+        ):
+            return None
 
         try:
             return file_class.objects.get(pk=file_id)

src/sentry/tasks/files.py

@@ -115,6 +115,10 @@ def copy_file_to_control_and_update_model(
 ):
     from sentry.models.files import ControlFile, File
 
+    if SiloMode.get_current_mode() != SiloMode.MONOLITH:
+        # We can only run this task in monolith mode.
+        return
+
     lock = f"copy-file-lock-{model_name}:{model_id}"
 
     with locks.get(lock, duration=60, name="copy-file-lock").acquire():

tests/sentry/api/endpoints/test_user_avatar.py

@@ -2,6 +2,8 @@ from base64 import b64encode
 from io import BytesIO
 from unittest import mock
 
+import pytest
+from django.test import override_settings
 from django.urls import reverse
 
 from sentry import options as options_store
@@ -9,6 +11,7 @@ from sentry.models.avatars.user_avatar import UserAvatar, UserAvatarType
 from sentry.models.files import ControlFile, File
 from sentry.silo.base import SiloMode
 from sentry.testutils.cases import APITestCase
+from sentry.testutils.hybrid_cloud import use_split_dbs
 from sentry.testutils.silo import assume_test_silo_mode, control_silo_test
 
 
@@ -161,6 +164,7 @@ class UserAvatarTest(APITestCase):
         assert isinstance(avatar.get_file(), ControlFile)
         assert ControlFile.objects.filter(id=avatar.control_file_id).exists()
 
+    @pytest.mark.skipif(use_split_dbs(), reason="requires single/monolith db mode")
     @mock.patch("sentry.tasks.files.copy_file_to_control_and_update_model.apply_async")
     def test_get_upload_use_control_during_update(self, mock_task):
         user = self.create_user(email="a@example.com")
@@ -168,12 +172,21 @@ class UserAvatarTest(APITestCase):
         with assume_test_silo_mode(SiloMode.REGION):
             photo = File.objects.create(name="test.png", type="avatar.file")
             photo.putfile(BytesIO(b"test"))
-        UserAvatar.objects.create(user=user, file_id=photo.id, avatar_type=UserAvatarType.UPLOAD)
+        avatar = UserAvatar.objects.create(
+            user=user, file_id=photo.id, avatar_type=UserAvatarType.UPLOAD
+        )
+        assert avatar
+        assert avatar.get_file_id()
+        assert avatar.file_id
+        assert not avatar.control_file_id
+        with assume_test_silo_mode(SiloMode.REGION):
+            assert isinstance(avatar.get_file(), File)
 
         self.login_as(user=user)
 
         url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"})
 
+        mock_task.reset_mock()
         # Run in monolith as this transitional operation can only occur there.
         with assume_test_silo_mode(SiloMode.MONOLITH):
             response = self.client.put(
@@ -198,6 +211,52 @@ class UserAvatarTest(APITestCase):
         assert avatar.control_file_id
         assert isinstance(avatar.get_file(), ControlFile)
 
+    @mock.patch("sentry.tasks.files.copy_file_to_control_and_update_model.apply_async")
+    def test_do_not_copy_file_to_control(self, mock_task):
+        user = self.create_user(email="a@example.com")
+
+        with assume_test_silo_mode(SiloMode.REGION):
+            photo = File.objects.create(name="test.png", type="avatar.file")
+            photo.putfile(BytesIO(b"test"))
+        avatar = UserAvatar.objects.create(
+            user=user, file_id=photo.id, avatar_type=UserAvatarType.UPLOAD
+        )
+
+        assert avatar
+        assert avatar.get_file_id()
+        assert avatar.file_id
+        assert not avatar.control_file_id
+        with assume_test_silo_mode(SiloMode.REGION):
+            assert isinstance(avatar.get_file(), File)
+
+        self.login_as(user=user)
+
+        url = reverse("sentry-api-0-user-avatar", kwargs={"user_id": "me"})
+
+        mock_task.reset_mock()
+        with assume_test_silo_mode(SiloMode.CONTROL), override_settings(SILO_MODE=SiloMode.CONTROL):
+            response = self.client.put(
+                url,
+                data={
+                    "avatar_type": "upload",
+                    "avatar_photo": b64encode(self.load_fixture("avatar.jpg")),
+                },
+                format="json",
+            )
+
+        assert response.status_code == 200, response.content
+        assert response.data["id"] == str(user.id)
+        assert response.data["avatar"]["avatarType"] == "upload"
+        assert response.data["avatar"]["avatarUuid"]
+        assert mock_task.call_count == 0
+
+        avatar = UserAvatar.objects.get(user=user)
+        assert avatar
+        assert avatar.get_file_id()
+        assert avatar.file_id is None, "non-control file relation be removed."
+        assert avatar.control_file_id
+        assert isinstance(avatar.get_file(), ControlFile)
+
     def test_put_upload_saves_to_control_file_with_separate_storage(self):
         with self.options(
             {