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