4
4
5
5
import re
6
6
import uuid
7
+ from unittest import mock
7
8
8
9
from django .core .files .storage import default_storage
9
10
from django .core .files .uploadedfile import SimpleUploadedFile
12
13
from rest_framework .test import APIClient
13
14
14
15
from core import factories
16
+ from core .api .viewsets import malware_detection
15
17
from core .tests .conftest import TEAM , USER , VIA
16
18
17
19
pytestmark = pytest .mark .django_db
@@ -59,7 +61,8 @@ def test_api_documents_attachment_upload_anonymous_success():
59
61
file = SimpleUploadedFile (name = "test.png" , content = PIXEL , content_type = "image/png" )
60
62
61
63
url = f"/api/v1.0/documents/{ document .id !s} /attachment-upload/"
62
- response = APIClient ().post (url , {"file" : file }, format = "multipart" )
64
+ with mock .patch .object (malware_detection , "analyse_file" ) as mock_analyse_file :
65
+ response = APIClient ().post (url , {"file" : file }, format = "multipart" )
63
66
64
67
assert response .status_code == 201
65
68
@@ -74,12 +77,13 @@ def test_api_documents_attachment_upload_anonymous_success():
74
77
assert document .attachments == [f"{ document .id !s} /attachments/{ file_id !s} .png" ]
75
78
76
79
# Now, check the metadata of the uploaded file
77
- key = file_path .replace ("/media" , "" )
80
+ key = file_path .replace ("/media/" , "" )
81
+ mock_analyse_file .assert_called_once_with (key , document_id = document .id )
78
82
file_head = default_storage .connection .meta .client .head_object (
79
83
Bucket = default_storage .bucket_name , Key = key
80
84
)
81
85
82
- assert file_head ["Metadata" ] == {"owner" : "None" }
86
+ assert file_head ["Metadata" ] == {"owner" : "None" , "status" : "processing" }
83
87
assert file_head ["ContentType" ] == "image/png"
84
88
assert file_head ["ContentDisposition" ] == 'inline; filename="test.png"'
85
89
@@ -139,14 +143,19 @@ def test_api_documents_attachment_upload_authenticated_success(reach, role):
139
143
file = SimpleUploadedFile (name = "test.png" , content = PIXEL , content_type = "image/png" )
140
144
141
145
url = f"/api/v1.0/documents/{ document .id !s} /attachment-upload/"
142
- response = client .post (url , {"file" : file }, format = "multipart" )
146
+ with mock .patch .object (malware_detection , "analyse_file" ) as mock_analyse_file :
147
+ response = client .post (url , {"file" : file }, format = "multipart" )
143
148
144
149
assert response .status_code == 201
145
150
146
151
pattern = re .compile (rf"^/media/{ document .id !s} /attachments/(.*)\.png" )
147
152
match = pattern .search (response .json ()["file" ])
148
153
file_id = match .group (1 )
149
154
155
+ mock_analyse_file .assert_called_once_with (
156
+ f"{ document .id !s} /attachments/{ file_id !s} .png" , document_id = document .id
157
+ )
158
+
150
159
# Validate that file_id is a valid UUID
151
160
uuid .UUID (file_id )
152
161
@@ -210,7 +219,8 @@ def test_api_documents_attachment_upload_success(via, role, mock_user_teams):
210
219
file = SimpleUploadedFile (name = "test.png" , content = PIXEL , content_type = "image/png" )
211
220
212
221
url = f"/api/v1.0/documents/{ document .id !s} /attachment-upload/"
213
- response = client .post (url , {"file" : file }, format = "multipart" )
222
+ with mock .patch .object (malware_detection , "analyse_file" ) as mock_analyse_file :
223
+ response = client .post (url , {"file" : file }, format = "multipart" )
214
224
215
225
assert response .status_code == 201
216
226
@@ -226,11 +236,12 @@ def test_api_documents_attachment_upload_success(via, role, mock_user_teams):
226
236
assert document .attachments == [f"{ document .id !s} /attachments/{ file_id !s} .png" ]
227
237
228
238
# Now, check the metadata of the uploaded file
229
- key = file_path .replace ("/media" , "" )
239
+ key = file_path .replace ("/media/" , "" )
240
+ mock_analyse_file .assert_called_once_with (key , document_id = document .id )
230
241
file_head = default_storage .connection .meta .client .head_object (
231
242
Bucket = default_storage .bucket_name , Key = key
232
243
)
233
- assert file_head ["Metadata" ] == {"owner" : str (user .id )}
244
+ assert file_head ["Metadata" ] == {"owner" : str (user .id ), "status" : "processing" }
234
245
assert file_head ["ContentType" ] == "image/png"
235
246
assert file_head ["ContentDisposition" ] == 'inline; filename="test.png"'
236
247
@@ -304,7 +315,8 @@ def test_api_documents_attachment_upload_fix_extension(
304
315
url = f"/api/v1.0/documents/{ document .id !s} /attachment-upload/"
305
316
306
317
file = SimpleUploadedFile (name = name , content = content )
307
- response = client .post (url , {"file" : file }, format = "multipart" )
318
+ with mock .patch .object (malware_detection , "analyse_file" ) as mock_analyse_file :
319
+ response = client .post (url , {"file" : file }, format = "multipart" )
308
320
309
321
assert response .status_code == 201
310
322
@@ -324,11 +336,16 @@ def test_api_documents_attachment_upload_fix_extension(
324
336
uuid .UUID (file_id )
325
337
326
338
# Now, check the metadata of the uploaded file
327
- key = file_path .replace ("/media" , "" )
339
+ key = file_path .replace ("/media/" , "" )
340
+ mock_analyse_file .assert_called_once_with (key , document_id = document .id )
328
341
file_head = default_storage .connection .meta .client .head_object (
329
342
Bucket = default_storage .bucket_name , Key = key
330
343
)
331
- assert file_head ["Metadata" ] == {"owner" : str (user .id ), "is_unsafe" : "true" }
344
+ assert file_head ["Metadata" ] == {
345
+ "owner" : str (user .id ),
346
+ "is_unsafe" : "true" ,
347
+ "status" : "processing" ,
348
+ }
332
349
assert file_head ["ContentType" ] == content_type
333
350
assert file_head ["ContentDisposition" ] == f'attachment; filename="{ name :s} "'
334
351
@@ -364,7 +381,8 @@ def test_api_documents_attachment_upload_unsafe():
364
381
file = SimpleUploadedFile (
365
382
name = "script.exe" , content = b"\x4d \x5a \x90 \x00 \x03 \x00 \x00 \x00 "
366
383
)
367
- response = client .post (url , {"file" : file }, format = "multipart" )
384
+ with mock .patch .object (malware_detection , "analyse_file" ) as mock_analyse_file :
385
+ response = client .post (url , {"file" : file }, format = "multipart" )
368
386
369
387
assert response .status_code == 201
370
388
@@ -381,11 +399,16 @@ def test_api_documents_attachment_upload_unsafe():
381
399
file_id = file_id .replace ("-unsafe" , "" )
382
400
uuid .UUID (file_id )
383
401
402
+ key = file_path .replace ("/media/" , "" )
403
+ mock_analyse_file .assert_called_once_with (key , document_id = document .id )
384
404
# Now, check the metadata of the uploaded file
385
- key = file_path .replace ("/media" , "" )
386
405
file_head = default_storage .connection .meta .client .head_object (
387
406
Bucket = default_storage .bucket_name , Key = key
388
407
)
389
- assert file_head ["Metadata" ] == {"owner" : str (user .id ), "is_unsafe" : "true" }
408
+ assert file_head ["Metadata" ] == {
409
+ "owner" : str (user .id ),
410
+ "is_unsafe" : "true" ,
411
+ "status" : "processing" ,
412
+ }
390
413
assert file_head ["ContentType" ] == "application/octet-stream"
391
414
assert file_head ["ContentDisposition" ] == 'attachment; filename="script.exe"'
0 commit comments