Coverage for tests/test_pdf_upload.py: 100%

1import io 

2from werkzeug.datastructures import FileStorage 

3import os 

4from app.models import User, Document 

5from app import db 

6 

7 

8 

9def login_user(client): 

10 test_user = User(username="testuser") 

11 test_user.set_password("password") 

12 db.session.add(test_user) 

13 db.session.commit() 

14 response = client.post( 

15 "/index", 

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

17 follow_redirects=True, 

18 ) 

19 assert response.status_code == 200 

20 

21 return client 

22 

23def test_upload_pdf_valid(client, app): 

24 app.config['WTF_CSRF_ENABLED'] = False 

25 

26 

27 with app.app_context(): 

28 client = login_user(client) 

29 # Path to the test file inside the test_data folder 

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

31 

32 # Open the file in binary read mode 

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

34 # Create a FileStorage object from the file content 

35 pdf_file = FileStorage( 

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

37 filename="test.pdf", 

38 content_type="application/pdf", 

39 ) 

40 

41 # Prepare the data dictionary with the file for upload 

42 data = { 

43 "pdf_file": (io.BytesIO(pdf_file.read()), "test.pdf"), 

44 } 

45 

46 

47 # Make the POST request to the upload route 

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

49 

50 # Get the JSON response using get_json() 

51 response_data = upload_response.get_json() 

52 

53 # Ensure the response data is not None before accessing 

54 assert response_data is not None # Ensure the response is valid JSON 

55 

56 # Check the JSON message 

57 assert upload_response.status_code == 200 

58 assert response_data["message"] == "File 'test.pdf' uploaded successfully!" 

59 

60 

61def test_upload_pdf_dup(client, app): 

62 app.config['WTF_CSRF_ENABLED'] = False 

63 

64 

65 with app.app_context(): 

66 client = login_user(client) 

67 # Path to the test file inside the test_data folder 

68 file_path = os.path.join(os.path.dirname(__file__), 'test_data', 'Dna.pdf') 

69 

70 # Open the file in binary read mode 

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

72 # Create a FileStorage object from the file content 

73 pdf_file = FileStorage( 

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

75 filename="Dna.pdf", 

76 content_type="application/pdf", 

77 ) 

78 

79 # Prepare the data dictionary with the file for upload 

80 data = { 

81 "pdf_file": (io.BytesIO(pdf_file.read()), "Dna.pdf"), 

82 } 

83 

84 # Make the POST request to the upload route First time 

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

86 

87 # Open the file again in binary read mode 

88 with open(file_path, 'rb') as f2: 

89 # Create a FileStorage object from the file content 

90 pdf_file2 = FileStorage( 

91 stream=io.BytesIO(f2.read()), 

92 filename="Dna.pdf", 

93 content_type="application/pdf", 

94 ) 

95 

96 # Prepare the data dictionary with the file for upload 

97 data = { 

98 "pdf_file": (io.BytesIO(pdf_file2.read()), "Dna.pdf"), 

99 } 

100 

101 

102 # Make the POST request to the upload route Second time 

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

104 

105 # Get the JSON response using get_json() 

106 response_data = upload_response_2.get_json() 

107 

108 # Ensure the response data is not None before accessing 

109 assert response_data is not None # Ensure the response is valid JSON 

110 

111 # Check the JSON message 

112 assert upload_response_2.status_code == 409 

113 assert response_data["error"] == "A document named 'Dna.pdf' already exists."