|
4 | 4 |
|
5 | 5 | from unittest.mock import AsyncMock, MagicMock, Mock, patch |
6 | 6 |
|
| 7 | +from azure.core.exceptions import ClientAuthenticationError, ResourceNotFoundError, ServiceRequestError |
7 | 8 | from openai import RateLimitError, BadRequestError, InternalServerError |
8 | 9 | import pytest |
9 | 10 | from flask.testing import FlaskClient |
@@ -923,3 +924,228 @@ def test_conversation_azure_byod_returns_correct_response_when_streaming_without |
923 | 924 | data |
924 | 925 | == '{"id": "response.id", "model": "mock-openai-model", "created": 0, "object": "response.object", "choices": [{"messages": [{"role": "assistant", "content": "mock content"}]}]}\n' |
925 | 926 | ) |
| 927 | + |
| 928 | + |
| 929 | +class TestGetFile: |
| 930 | + """Test the get_file endpoint for downloading files from blob storage.""" |
| 931 | + |
| 932 | + @patch("create_app.AzureBlobStorageClient") |
| 933 | + def test_get_file_success(self, mock_blob_client_class, client): |
| 934 | + """Test successful file download with proper headers.""" |
| 935 | + # given |
| 936 | + filename = "test_document.pdf" |
| 937 | + file_content = b"Mock file content for PDF document" |
| 938 | + |
| 939 | + mock_blob_client = MagicMock() |
| 940 | + mock_blob_client_class.return_value = mock_blob_client |
| 941 | + mock_blob_client.file_exists.return_value = True |
| 942 | + mock_blob_client.download_file.return_value = file_content |
| 943 | + |
| 944 | + # when |
| 945 | + response = client.get(f"/api/files/{filename}") |
| 946 | + |
| 947 | + # then |
| 948 | + assert response.status_code == 200 |
| 949 | + assert response.data == file_content |
| 950 | + assert response.headers["Content-Type"] == "application/pdf" |
| 951 | + assert response.headers["Content-Disposition"] == f'inline; filename="{filename}"' |
| 952 | + assert response.headers["Content-Length"] == str(len(file_content)) |
| 953 | + assert response.headers["Cache-Control"] == "public, max-age=3600" |
| 954 | + assert response.headers["X-Content-Type-Options"] == "nosniff" |
| 955 | + assert response.headers["X-Frame-Options"] == "DENY" |
| 956 | + assert response.headers["Content-Security-Policy"] == "default-src 'none'" |
| 957 | + |
| 958 | + # Verify blob client was initialized with correct container |
| 959 | + mock_blob_client_class.assert_called_once_with(container_name="documents") |
| 960 | + mock_blob_client.file_exists.assert_called_once_with(filename) |
| 961 | + mock_blob_client.download_file.assert_called_once_with(filename) |
| 962 | + |
| 963 | + @patch("create_app.AzureBlobStorageClient") |
| 964 | + def test_get_file_with_unknown_mime_type(self, mock_blob_client_class, client): |
| 965 | + """Test file download with unknown file extension.""" |
| 966 | + # given |
| 967 | + filename = "test_file.unknownext" |
| 968 | + file_content = b"Mock file content" |
| 969 | + |
| 970 | + mock_blob_client = MagicMock() |
| 971 | + mock_blob_client_class.return_value = mock_blob_client |
| 972 | + mock_blob_client.file_exists.return_value = True |
| 973 | + mock_blob_client.download_file.return_value = file_content |
| 974 | + |
| 975 | + # when |
| 976 | + response = client.get(f"/api/files/{filename}") |
| 977 | + |
| 978 | + # then |
| 979 | + assert response.status_code == 200 |
| 980 | + assert response.headers["Content-Type"] == "application/octet-stream" |
| 981 | + |
| 982 | + @patch("create_app.AzureBlobStorageClient") |
| 983 | + def test_get_file_large_file_warning(self, mock_blob_client_class, client): |
| 984 | + """Test that large files are handled properly with logging.""" |
| 985 | + # given |
| 986 | + filename = "large_document.pdf" |
| 987 | + file_content = b"x" * (11 * 1024 * 1024) # 11MB file |
| 988 | + |
| 989 | + mock_blob_client = MagicMock() |
| 990 | + mock_blob_client_class.return_value = mock_blob_client |
| 991 | + mock_blob_client.file_exists.return_value = True |
| 992 | + mock_blob_client.download_file.return_value = file_content |
| 993 | + |
| 994 | + # when |
| 995 | + response = client.get(f"/api/files/{filename}") |
| 996 | + |
| 997 | + # then |
| 998 | + assert response.status_code == 200 |
| 999 | + assert len(response.data) == len(file_content) |
| 1000 | + |
| 1001 | + def test_get_file_empty_filename(self, client): |
| 1002 | + """Test error response when filename is empty.""" |
| 1003 | + # when |
| 1004 | + response = client.get("/api/files/") |
| 1005 | + |
| 1006 | + # then |
| 1007 | + # This should result in a 404 as the route won't match |
| 1008 | + assert response.status_code == 404 |
| 1009 | + |
| 1010 | + def test_get_file_invalid_filename_too_long(self, client): |
| 1011 | + """Test error response for filenames that are too long.""" |
| 1012 | + # given |
| 1013 | + filename = "a" * 256 # 256 characters, exceeds 255 limit |
| 1014 | + |
| 1015 | + # when |
| 1016 | + response = client.get(f"/api/files/{filename}") |
| 1017 | + |
| 1018 | + # then |
| 1019 | + assert response.status_code == 400 |
| 1020 | + assert response.json == {"error": "Filename too long"} |
| 1021 | + |
| 1022 | + @patch("create_app.AzureBlobStorageClient") |
| 1023 | + def test_get_file_not_exists_in_storage(self, mock_blob_client_class, client): |
| 1024 | + """Test error response when file doesn't exist in blob storage.""" |
| 1025 | + # given |
| 1026 | + filename = "nonexistent.pdf" |
| 1027 | + |
| 1028 | + mock_blob_client = MagicMock() |
| 1029 | + mock_blob_client_class.return_value = mock_blob_client |
| 1030 | + mock_blob_client.file_exists.return_value = False |
| 1031 | + |
| 1032 | + # when |
| 1033 | + response = client.get(f"/api/files/{filename}") |
| 1034 | + |
| 1035 | + # then |
| 1036 | + assert response.status_code == 404 |
| 1037 | + assert response.json == {"error": "File not found"} |
| 1038 | + mock_blob_client.file_exists.assert_called_once_with(filename) |
| 1039 | + mock_blob_client.download_file.assert_not_called() |
| 1040 | + |
| 1041 | + @patch("create_app.AzureBlobStorageClient") |
| 1042 | + def test_get_file_client_authentication_error(self, mock_blob_client_class, client): |
| 1043 | + """Test handling of Azure ClientAuthenticationError.""" |
| 1044 | + # given |
| 1045 | + filename = "test.pdf" |
| 1046 | + |
| 1047 | + mock_blob_client = MagicMock() |
| 1048 | + mock_blob_client_class.return_value = mock_blob_client |
| 1049 | + mock_blob_client.file_exists.side_effect = ClientAuthenticationError("Auth failed") |
| 1050 | + |
| 1051 | + # when |
| 1052 | + response = client.get(f"/api/files/{filename}") |
| 1053 | + |
| 1054 | + # then |
| 1055 | + assert response.status_code == 401 |
| 1056 | + assert response.json == {"error": "Authentication failed"} |
| 1057 | + |
| 1058 | + @patch("create_app.AzureBlobStorageClient") |
| 1059 | + def test_get_file_resource_not_found_error(self, mock_blob_client_class, client): |
| 1060 | + """Test handling of Azure ResourceNotFoundError.""" |
| 1061 | + # given |
| 1062 | + filename = "test.pdf" |
| 1063 | + |
| 1064 | + mock_blob_client = MagicMock() |
| 1065 | + mock_blob_client_class.return_value = mock_blob_client |
| 1066 | + mock_blob_client.file_exists.side_effect = ResourceNotFoundError("Resource not found") |
| 1067 | + |
| 1068 | + # when |
| 1069 | + response = client.get(f"/api/files/{filename}") |
| 1070 | + |
| 1071 | + # then |
| 1072 | + assert response.status_code == 404 |
| 1073 | + assert response.json == {"error": "File not found"} |
| 1074 | + |
| 1075 | + @patch("create_app.AzureBlobStorageClient") |
| 1076 | + def test_get_file_service_request_error(self, mock_blob_client_class, client): |
| 1077 | + """Test handling of Azure ServiceRequestError.""" |
| 1078 | + # given |
| 1079 | + filename = "test.pdf" |
| 1080 | + |
| 1081 | + mock_blob_client = MagicMock() |
| 1082 | + mock_blob_client_class.return_value = mock_blob_client |
| 1083 | + mock_blob_client.file_exists.side_effect = ServiceRequestError("Service unavailable") |
| 1084 | + |
| 1085 | + # when |
| 1086 | + response = client.get(f"/api/files/{filename}") |
| 1087 | + |
| 1088 | + # then |
| 1089 | + assert response.status_code == 503 |
| 1090 | + assert response.json == {"error": "Storage service unavailable"} |
| 1091 | + |
| 1092 | + @patch("create_app.AzureBlobStorageClient") |
| 1093 | + def test_get_file_unexpected_exception(self, mock_blob_client_class, client): |
| 1094 | + """Test handling of unexpected exceptions.""" |
| 1095 | + # given |
| 1096 | + filename = "test.pdf" |
| 1097 | + |
| 1098 | + mock_blob_client = MagicMock() |
| 1099 | + mock_blob_client_class.return_value = mock_blob_client |
| 1100 | + mock_blob_client.file_exists.side_effect = Exception("Unexpected error") |
| 1101 | + |
| 1102 | + # when |
| 1103 | + response = client.get(f"/api/files/{filename}") |
| 1104 | + |
| 1105 | + # then |
| 1106 | + assert response.status_code == 500 |
| 1107 | + assert response.json == {"error": "Internal server error"} |
| 1108 | + |
| 1109 | + @patch("create_app.AzureBlobStorageClient") |
| 1110 | + def test_get_file_download_exception(self, mock_blob_client_class, client): |
| 1111 | + """Test handling of exceptions during file download.""" |
| 1112 | + # given |
| 1113 | + filename = "test.pdf" |
| 1114 | + |
| 1115 | + mock_blob_client = MagicMock() |
| 1116 | + mock_blob_client_class.return_value = mock_blob_client |
| 1117 | + mock_blob_client.file_exists.return_value = True |
| 1118 | + mock_blob_client.download_file.side_effect = Exception("Download failed") |
| 1119 | + |
| 1120 | + # when |
| 1121 | + response = client.get(f"/api/files/{filename}") |
| 1122 | + |
| 1123 | + # then |
| 1124 | + assert response.status_code == 500 |
| 1125 | + assert response.json == {"error": "Internal server error"} |
| 1126 | + |
| 1127 | + def test_get_file_valid_filenames(self, client): |
| 1128 | + """Test that valid filenames with allowed characters pass validation.""" |
| 1129 | + # Mock the blob client to avoid actual Azure calls |
| 1130 | + with patch("create_app.AzureBlobStorageClient") as mock_blob_client_class: |
| 1131 | + mock_blob_client = MagicMock() |
| 1132 | + mock_blob_client_class.return_value = mock_blob_client |
| 1133 | + mock_blob_client.file_exists.return_value = True |
| 1134 | + mock_blob_client.download_file.return_value = b"test content" |
| 1135 | + |
| 1136 | + valid_filenames = [ |
| 1137 | + "document.pdf", |
| 1138 | + "file_name.txt", |
| 1139 | + "file-name.docx", |
| 1140 | + "file name.xlsx", |
| 1141 | + "test123.json", |
| 1142 | + "a.b", |
| 1143 | + "very_long_but_valid_filename_with_underscores.pdf" |
| 1144 | + ] |
| 1145 | + |
| 1146 | + for filename in valid_filenames: |
| 1147 | + # when |
| 1148 | + response = client.get(f"/api/files/{filename}") |
| 1149 | + |
| 1150 | + # then |
| 1151 | + assert response.status_code == 200, f"Failed for filename: {filename}" |
0 commit comments