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),
],
)
|
