Prevent local quarantined media from being claimed by media retention (#12972)

1 files changed, 96 insertions, 13 deletions

diff --git a/tests/rest/media/test_media_retention.py b/tests/rest/media/test_media_retention.py
index b98a5cd586..14af07c5af 100644
--- a/tests/rest/media/test_media_retention.py
+++ b/tests/rest/media/test_media_retention.py
@@ -53,13 +53,16 @@ class MediaRetentionTestCase(unittest.HomeserverTestCase):
         # Create a user to upload media with
         test_user_id = self.register_user("alice", "password")
 
-        # Inject media (3 images each; recently accessed, old access, never accessed)
-        # into both the local store and the remote cache
+        # Inject media (recently accessed, old access, never accessed, old access
+        # quarantined media) into both the local store and the remote cache, plus
+        # one additional local media that is marked as protected from quarantine.
         media_repository = hs.get_media_repository()
         test_media_content = b"example string"
 
-        def _create_media_and_set_last_accessed(
+        def _create_media_and_set_attributes(
             last_accessed_ms: Optional[int],
+            is_quarantined: Optional[bool] = False,
+            is_protected: Optional[bool] = False,
         ) -> str:
             # "Upload" some media to the local media store
             mxc_uri = self.get_success(
@@ -84,10 +87,31 @@ class MediaRetentionTestCase(unittest.HomeserverTestCase):
                     )
                 )
 
+            if is_quarantined:
+                # Mark this media as quarantined
+                self.get_success(
+                    self.store.quarantine_media_by_id(
+                        server_name=self.hs.config.server.server_name,
+                        media_id=media_id,
+                        quarantined_by="@theadmin:test",
+                    )
+                )
+
+            if is_protected:
+                # Mark this media as protected from quarantine
+                self.get_success(
+                    self.store.mark_local_media_as_safe(
+                        media_id=media_id,
+                        safe=True,
+                    )
+                )
+
             return media_id
 
-        def _cache_remote_media_and_set_last_accessed(
-            media_id: str, last_accessed_ms: Optional[int]
+        def _cache_remote_media_and_set_attributes(
+            media_id: str,
+            last_accessed_ms: Optional[int],
+            is_quarantined: Optional[bool] = False,
         ) -> str:
             # Pretend to cache some remote media
             self.get_success(
@@ -112,23 +136,58 @@ class MediaRetentionTestCase(unittest.HomeserverTestCase):
                     )
                 )
 
+            if is_quarantined:
+                # Mark this media as quarantined
+                self.get_success(
+                    self.store.quarantine_media_by_id(
+                        server_name=self.remote_server_name,
+                        media_id=media_id,
+                        quarantined_by="@theadmin:test",
+                    )
+                )
+
             return media_id
 
         # Start with the local media store
-        self.local_recently_accessed_media = _create_media_and_set_last_accessed(
-            self.THIRTY_DAYS_IN_MS
+        self.local_recently_accessed_media = _create_media_and_set_attributes(
+            last_accessed_ms=self.THIRTY_DAYS_IN_MS,
         )
-        self.local_not_recently_accessed_media = _create_media_and_set_last_accessed(
-            self.ONE_DAY_IN_MS
+        self.local_not_recently_accessed_media = _create_media_and_set_attributes(
+            last_accessed_ms=self.ONE_DAY_IN_MS,
+        )
+        self.local_not_recently_accessed_quarantined_media = (
+            _create_media_and_set_attributes(
+                last_accessed_ms=self.ONE_DAY_IN_MS,
+                is_quarantined=True,
+            )
+        )
+        self.local_not_recently_accessed_protected_media = (
+            _create_media_and_set_attributes(
+                last_accessed_ms=self.ONE_DAY_IN_MS,
+                is_protected=True,
+            )
+        )
+        self.local_never_accessed_media = _create_media_and_set_attributes(
+            last_accessed_ms=None,
         )
-        self.local_never_accessed_media = _create_media_and_set_last_accessed(None)
 
         # And now the remote media store
-        self.remote_recently_accessed_media = _cache_remote_media_and_set_last_accessed(
-            "a", self.THIRTY_DAYS_IN_MS
+        self.remote_recently_accessed_media = _cache_remote_media_and_set_attributes(
+            media_id="a",
+            last_accessed_ms=self.THIRTY_DAYS_IN_MS,
         )
         self.remote_not_recently_accessed_media = (
-            _cache_remote_media_and_set_last_accessed("b", self.ONE_DAY_IN_MS)
+            _cache_remote_media_and_set_attributes(
+                media_id="b",
+                last_accessed_ms=self.ONE_DAY_IN_MS,
+            )
+        )
+        self.remote_not_recently_accessed_quarantined_media = (
+            _cache_remote_media_and_set_attributes(
+                media_id="c",
+                last_accessed_ms=self.ONE_DAY_IN_MS,
+                is_quarantined=True,
+            )
         )
         # Remote media will always have a "last accessed" attribute, as it would not
         # be fetched from the remote homeserver unless instigated by a user.
@@ -163,8 +222,20 @@ class MediaRetentionTestCase(unittest.HomeserverTestCase):
             ],
             not_purged=[
                 (self.hs.config.server.server_name, self.local_recently_accessed_media),
+                (
+                    self.hs.config.server.server_name,
+                    self.local_not_recently_accessed_quarantined_media,
+                ),
+                (
+                    self.hs.config.server.server_name,
+                    self.local_not_recently_accessed_protected_media,
+                ),
                 (self.remote_server_name, self.remote_recently_accessed_media),
                 (self.remote_server_name, self.remote_not_recently_accessed_media),
+                (
+                    self.remote_server_name,
+                    self.remote_not_recently_accessed_quarantined_media,
+                ),
             ],
         )
 
@@ -199,6 +270,18 @@ class MediaRetentionTestCase(unittest.HomeserverTestCase):
                     self.hs.config.server.server_name,
                     self.local_not_recently_accessed_media,
                 ),
+                (
+                    self.hs.config.server.server_name,
+                    self.local_not_recently_accessed_quarantined_media,
+                ),
+                (
+                    self.hs.config.server.server_name,
+                    self.local_not_recently_accessed_protected_media,
+                ),
+                (
+                    self.remote_server_name,
+                    self.remote_not_recently_accessed_quarantined_media,
+                ),
                 (self.hs.config.server.server_name, self.local_never_accessed_media),
             ],
         )