From 0a296b2a6e0ac871e03635e49c09ce6faa6f073b Mon Sep 17 00:00:00 2001
From: Alex Vandiver <alexmv@zulip.com>
Date: Thu, 20 Jun 2024 16:54:49 +0000
Subject: [PATCH] upload: Start storing content-type for new uploads.

---
 analytics/tests/test_counts.py     | 14 ++++++++++----
 zerver/lib/upload/__init__.py      |  9 ++++++++-
 zerver/tests/test_message_fetch.py |  4 +++-
 zerver/tests/test_retention.py     |  8 ++++++--
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/analytics/tests/test_counts.py b/analytics/tests/test_counts.py
index f4c9c7a715..0d9b9b37ad 100644
--- a/analytics/tests/test_counts.py
+++ b/analytics/tests/test_counts.py
@@ -204,7 +204,12 @@ class AnalyticsTestCase(ZulipTestCase):
         return Message.objects.create(**kwargs)
 
     def create_attachment(
-        self, user_profile: UserProfile, filename: str, size: int, create_time: datetime
+        self,
+        user_profile: UserProfile,
+        filename: str,
+        size: int,
+        create_time: datetime,
+        content_type: str,
     ) -> Attachment:
         return Attachment.objects.create(
             file_name=filename,
@@ -213,6 +218,7 @@ class AnalyticsTestCase(ZulipTestCase):
             realm=user_profile.realm,
             size=size,
             create_time=create_time,
+            content_type=content_type,
         )
 
     # kwargs should only ever be a UserProfile or Stream.
@@ -551,9 +557,9 @@ class TestCountStats(AnalyticsTestCase):
         user2 = self.create_user()
         user_second_realm = self.create_user(realm=self.second_realm)
 
-        self.create_attachment(user1, "file1", 100, self.TIME_LAST_HOUR)
-        attachment2 = self.create_attachment(user2, "file2", 200, self.TIME_LAST_HOUR)
-        self.create_attachment(user_second_realm, "file3", 10, self.TIME_LAST_HOUR)
+        self.create_attachment(user1, "file1", 100, self.TIME_LAST_HOUR, "text/plain")
+        attachment2 = self.create_attachment(user2, "file2", 200, self.TIME_LAST_HOUR, "text/plain")
+        self.create_attachment(user_second_realm, "file3", 10, self.TIME_LAST_HOUR, "text/plain")
 
         do_fill_count_stat_at_hour(stat, self.TIME_ZERO)
 
diff --git a/zerver/lib/upload/__init__.py b/zerver/lib/upload/__init__.py
index b8ee8cd1d0..a0ecccbdb0 100644
--- a/zerver/lib/upload/__init__.py
+++ b/zerver/lib/upload/__init__.py
@@ -41,7 +41,12 @@ def check_upload_within_quota(realm: Realm, uploaded_file_size: int) -> None:
 
 
 def create_attachment(
-    file_name: str, path_id: str, user_profile: UserProfile, realm: Realm, file_size: int
+    file_name: str,
+    path_id: str,
+    user_profile: UserProfile,
+    realm: Realm,
+    file_size: int,
+    content_type: str,
 ) -> None:
     assert (user_profile.realm_id == realm.id) or is_cross_realm_bot_email(
         user_profile.delivery_email
@@ -52,6 +57,7 @@ def create_attachment(
         owner=user_profile,
         realm=realm,
         size=file_size,
+        content_type=content_type,
     )
     from zerver.actions.uploads import notify_attachment_update
 
@@ -143,6 +149,7 @@ def upload_message_attachment(
         user_profile,
         target_realm,
         uploaded_file_size,
+        content_type,
     )
     return f"/user_uploads/{path_id}"
 
diff --git a/zerver/tests/test_message_fetch.py b/zerver/tests/test_message_fetch.py
index 593a975601..fe3b772344 100644
--- a/zerver/tests/test_message_fetch.py
+++ b/zerver/tests/test_message_fetch.py
@@ -4526,7 +4526,9 @@ class MessageHasKeywordsTest(ZulipTestCase):
         ]
 
         for file_name, path_id, size in dummy_files:
-            create_attachment(file_name, path_id, user_profile, user_profile.realm, size)
+            create_attachment(
+                file_name, path_id, user_profile, user_profile.realm, size, "text/plain"
+            )
 
         # return path ids
         return [x[1] for x in dummy_files]
diff --git a/zerver/tests/test_retention.py b/zerver/tests/test_retention.py
index 6266c48715..36e0196f45 100644
--- a/zerver/tests/test_retention.py
+++ b/zerver/tests/test_retention.py
@@ -190,7 +190,9 @@ class ArchiveMessagesTestingBase(RetentionTestingBase):
         ]
 
         for file_name, path_id, size in dummy_files:
-            create_attachment(file_name, path_id, user_profile, user_profile.realm, size)
+            create_attachment(
+                file_name, path_id, user_profile, user_profile.realm, size, "text/plain"
+            )
 
         self.subscribe(user_profile, "Denmark")
         body = (
@@ -603,7 +605,9 @@ class MoveMessageToArchiveBase(RetentionTestingBase):
         ]
         user_profile = self.example_user("hamlet")
         for file_name, path_id, size in dummy_files:
-            create_attachment(file_name, path_id, user_profile, user_profile.realm, size)
+            create_attachment(
+                file_name, path_id, user_profile, user_profile.realm, size, "text/plain"
+            )
 
     def _assert_archive_empty(self) -> None:
         self.assertFalse(ArchivedUserMessage.objects.exists())