diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/rest/media/test_media_retention.py | 109 |
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), ], ) |