Coverage for tests/test_download.py: 100%

1import os 

2import io 

3from werkzeug.datastructures import FileStorage 

4from app.models import Document, User 

5from app import db 

6from sqlalchemy import select 

7 

8def login_user(client): 

9 test_user = User(username="testuser") 

10 test_user.set_password("password") 

11 db.session.add(test_user) 

12 db.session.commit() 

13 response = client.post( 

14 "/index", 

15 data={"username": "testuser", "password": "password"}, 

16 follow_redirects=True, 

17 ) 

18 assert response.status_code == 200 

19 

20 return client 

21 

22 

23def upload_test_document(client, app): 

24 """ 

25 Helper function to upload a test document. 

26 Returns the uploaded document object. 

27 """ 

28 app.config['WTF_CSRF_ENABLED'] = False 

29 

30 with app.app_context(): 

31 client = login_user(client) 

32 # Path to the test file inside the test_data folder 

33 file_path = os.path.join(os.path.dirname(__file__), 'test_data', 'test.pdf') 

34 

35 # Open the file in binary read mode and create a FileStorage object 

36 with open(file_path, 'rb') as f: 

37 pdf_file = FileStorage( 

38 stream=io.BytesIO(f.read()), 

39 filename="test.pdf", 

40 content_type="application/pdf", 

41 ) 

42 

43 # Prepare the data dictionary with the file for upload 

44 data = {"pdf_file": (pdf_file, "test.pdf")} 

45 

46 # Make the POST request to the upload route 

47 response = client.post("/admin", data=data, content_type='multipart/form-data') 

48 assert response.status_code == 200, f"Upload failed: {response.json}" 

49 assert "document" in response.json 

50 

51 

52 return db.session.query(Document).filter_by(document_name="test").first() 

53 

54 

55def test_download_valid_document(client, app): 

56 """ 

57 Test that a pdf document can be downloaded. 

58 """ 

59 

60 uploaded_doc = upload_test_document(client, app) 

61 item_id = uploaded_doc.id 

62 

63 response = client.get(f"/download/{item_id}") 

64 

65 

66 with app.app_context(): 

67 

68 

69 assert response.status_code == 200 

70 assert response.mimetype == "application/pdf" 

71 assert response.headers["Content-Disposition"].startswith("attachment;") 

72 assert f"{uploaded_doc.document_name}.{uploaded_doc.document_type}" in response.headers["Content-Disposition"] 

73 assert isinstance(response.data, bytes) 

74 assert len(response.data) > 0 

75 

76def test_download_non_existent_document(client, app): 

77 """ 

78 Test that a pdf document that does not exist returns an error. 

79 """ 

80 non_existent_id = 9999 

81 

82 response = client.get(f"/download/{non_existent_id}") 

83 

84 assert response.status_code == 404 

85 assert not response.json["success"], "Response should indicate failure" 

86 assert response.json["message"] == f"Item {non_existent_id} not found" 

87 

88def test_download_database_error(client, app, monkeypatch): 

89 """ 

90 Simulates an internal DB error during download to ensure proper 500 response. 

91 Expected: 500 Internal Server Error 

92 """ 

93 

94 def faulty_get(*args, **kwargs): 

95 raise Exception("Simulated DB failure") 

96 

97 

98 uploaded_doc = upload_test_document(client, app) 

99 assert uploaded_doc is not None, "Failed to retrieve uploaded document" 

100 

101 item_id = uploaded_doc.id 

102 

103 with app.app_context(): 

104 

105 monkeypatch.setattr(db.session, "query", lambda *a, **k: type("MockQuery", (), {"get": faulty_get})()) 

106 

107 response = client.get(f"/download/{item_id}") 

108 

109 assert response.status_code == 500, f"Expected 500, but got {response.status_code}" 

110 assert not response.json["success"], "Response should indicate failure" 

111 assert response.json["message"] == "Failed to download document" 

112