@@ -53,13 +53,16 @@ def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
53
53
# Create a user to upload media with
54
54
test_user_id = self .register_user ("alice" , "password" )
55
55
56
- # Inject media (3 images each; recently accessed, old access, never accessed)
57
- # into both the local store and the remote cache
56
+ # Inject media (recently accessed, old access, never accessed, old access
57
+ # quarantined media) into both the local store and the remote cache, plus
58
+ # one additional local media that is marked as protected from quarantine.
58
59
media_repository = hs .get_media_repository ()
59
60
test_media_content = b"example string"
60
61
61
- def _create_media_and_set_last_accessed (
62
+ def _create_media_and_set_attributes (
62
63
last_accessed_ms : Optional [int ],
64
+ is_quarantined : Optional [bool ] = False ,
65
+ is_protected : Optional [bool ] = False ,
63
66
) -> str :
64
67
# "Upload" some media to the local media store
65
68
mxc_uri = self .get_success (
@@ -84,10 +87,31 @@ def _create_media_and_set_last_accessed(
84
87
)
85
88
)
86
89
90
+ if is_quarantined :
91
+ # Mark this media as quarantined
92
+ self .get_success (
93
+ self .store .quarantine_media_by_id (
94
+ server_name = self .hs .config .server .server_name ,
95
+ media_id = media_id ,
96
+ quarantined_by = "@theadmin:test" ,
97
+ )
98
+ )
99
+
100
+ if is_protected :
101
+ # Mark this media as protected from quarantine
102
+ self .get_success (
103
+ self .store .mark_local_media_as_safe (
104
+ media_id = media_id ,
105
+ safe = True ,
106
+ )
107
+ )
108
+
87
109
return media_id
88
110
89
- def _cache_remote_media_and_set_last_accessed (
90
- media_id : str , last_accessed_ms : Optional [int ]
111
+ def _cache_remote_media_and_set_attributes (
112
+ media_id : str ,
113
+ last_accessed_ms : Optional [int ],
114
+ is_quarantined : Optional [bool ] = False ,
91
115
) -> str :
92
116
# Pretend to cache some remote media
93
117
self .get_success (
@@ -112,23 +136,58 @@ def _cache_remote_media_and_set_last_accessed(
112
136
)
113
137
)
114
138
139
+ if is_quarantined :
140
+ # Mark this media as quarantined
141
+ self .get_success (
142
+ self .store .quarantine_media_by_id (
143
+ server_name = self .remote_server_name ,
144
+ media_id = media_id ,
145
+ quarantined_by = "@theadmin:test" ,
146
+ )
147
+ )
148
+
115
149
return media_id
116
150
117
151
# Start with the local media store
118
- self .local_recently_accessed_media = _create_media_and_set_last_accessed (
119
- self .THIRTY_DAYS_IN_MS
152
+ self .local_recently_accessed_media = _create_media_and_set_attributes (
153
+ last_accessed_ms = self .THIRTY_DAYS_IN_MS ,
120
154
)
121
- self .local_not_recently_accessed_media = _create_media_and_set_last_accessed (
122
- self .ONE_DAY_IN_MS
155
+ self .local_not_recently_accessed_media = _create_media_and_set_attributes (
156
+ last_accessed_ms = self .ONE_DAY_IN_MS ,
157
+ )
158
+ self .local_not_recently_accessed_quarantined_media = (
159
+ _create_media_and_set_attributes (
160
+ last_accessed_ms = self .ONE_DAY_IN_MS ,
161
+ is_quarantined = True ,
162
+ )
163
+ )
164
+ self .local_not_recently_accessed_protected_media = (
165
+ _create_media_and_set_attributes (
166
+ last_accessed_ms = self .ONE_DAY_IN_MS ,
167
+ is_protected = True ,
168
+ )
169
+ )
170
+ self .local_never_accessed_media = _create_media_and_set_attributes (
171
+ last_accessed_ms = None ,
123
172
)
124
- self .local_never_accessed_media = _create_media_and_set_last_accessed (None )
125
173
126
174
# And now the remote media store
127
- self .remote_recently_accessed_media = _cache_remote_media_and_set_last_accessed (
128
- "a" , self .THIRTY_DAYS_IN_MS
175
+ self .remote_recently_accessed_media = _cache_remote_media_and_set_attributes (
176
+ media_id = "a" ,
177
+ last_accessed_ms = self .THIRTY_DAYS_IN_MS ,
129
178
)
130
179
self .remote_not_recently_accessed_media = (
131
- _cache_remote_media_and_set_last_accessed ("b" , self .ONE_DAY_IN_MS )
180
+ _cache_remote_media_and_set_attributes (
181
+ media_id = "b" ,
182
+ last_accessed_ms = self .ONE_DAY_IN_MS ,
183
+ )
184
+ )
185
+ self .remote_not_recently_accessed_quarantined_media = (
186
+ _cache_remote_media_and_set_attributes (
187
+ media_id = "c" ,
188
+ last_accessed_ms = self .ONE_DAY_IN_MS ,
189
+ is_quarantined = True ,
190
+ )
132
191
)
133
192
# Remote media will always have a "last accessed" attribute, as it would not
134
193
# be fetched from the remote homeserver unless instigated by a user.
@@ -163,8 +222,20 @@ def test_local_media_retention(self) -> None:
163
222
],
164
223
not_purged = [
165
224
(self .hs .config .server .server_name , self .local_recently_accessed_media ),
225
+ (
226
+ self .hs .config .server .server_name ,
227
+ self .local_not_recently_accessed_quarantined_media ,
228
+ ),
229
+ (
230
+ self .hs .config .server .server_name ,
231
+ self .local_not_recently_accessed_protected_media ,
232
+ ),
166
233
(self .remote_server_name , self .remote_recently_accessed_media ),
167
234
(self .remote_server_name , self .remote_not_recently_accessed_media ),
235
+ (
236
+ self .remote_server_name ,
237
+ self .remote_not_recently_accessed_quarantined_media ,
238
+ ),
168
239
],
169
240
)
170
241
@@ -199,6 +270,18 @@ def test_remote_media_cache_retention(self) -> None:
199
270
self .hs .config .server .server_name ,
200
271
self .local_not_recently_accessed_media ,
201
272
),
273
+ (
274
+ self .hs .config .server .server_name ,
275
+ self .local_not_recently_accessed_quarantined_media ,
276
+ ),
277
+ (
278
+ self .hs .config .server .server_name ,
279
+ self .local_not_recently_accessed_protected_media ,
280
+ ),
281
+ (
282
+ self .remote_server_name ,
283
+ self .remote_not_recently_accessed_quarantined_media ,
284
+ ),
202
285
(self .hs .config .server .server_name , self .local_never_accessed_media ),
203
286
],
204
287
)
0 commit comments