Merge pull request #223 from taketo1113/api-upload

Add Upload Reports API

Author: mehaase

src/tram/urls.py

@@ -41,12 +41,13 @@
     path("api/", include(router.urls)),
     path("api/token/", TokenObtainPairView.as_view(), name="token_obtain_pair"),
     path("api/token/refresh/", TokenRefreshView.as_view(), name="token_refresh"),
+    path("api/upload/", views.upload_api),
     path("api/download/<int:doc_id>", views.download_document),
     path("api/train-model/<name>", views.train_model),
     path("docs/", TemplateView.as_view(template_name="tram_documentation.html")),
     path("login/", auth_views.LoginView.as_view()),
     path("logout/", auth_views.LogoutView.as_view()),
-    path("upload/", views.upload),
+    path("upload/", views.upload_web),
     path("admin/", admin.site.urls),
     path("ml/", views.ml_home),
     path("ml/techniques/<str:attack_id>", views.ml_technique_sentences),

src/tram/views.py

@@ -130,6 +130,15 @@ def index(request):
 
 @login_required
 @require_POST
+def upload_web(request):
+    return upload(request)
+
+
+@api_view(["POST"])
+def upload_api(request):
+    return upload(request)
+
+
 def upload(request):
     """Places a file into ml-pipeline for analysis"""
     # Initialize the processing job.
@@ -138,6 +147,9 @@ def upload(request):
     # Initialize response.
     response = {"message": "File saved for processing."}
 
+    if request.FILES.get("file") is None:
+        return HttpResponseBadRequest("'file' field is required but was not provided")
+
     file_content_type = request.FILES["file"].content_type
     if file_content_type in (
         "application/pdf",  # .pdf files
@@ -165,7 +177,6 @@ def upload(request):
 
     return JsonResponse(response)
 
-
 @login_required
 def ml_home(request):
     techniques = AttackObject.get_sentence_counts()

tests/tram/test_views.py

@@ -180,6 +180,29 @@ def test_upload_unsupported_file_type_causes_bad_request(self, logged_in_client)
         assert response.content == b"Unsupported file type"
 
 
+@pytest.mark.django_db
+class TestUploadApi:
+    def test_upload_api(self, logged_in_client):
+        # Arrange
+        f = SimpleUploadedFile(
+            "test-report.pdf", b"test file content", content_type="application/pdf"
+        )
+        data = {"file": f}
+
+        # Act
+        response = logged_in_client.post("/api/upload/", data)
+
+        # Assert
+        assert response.status_code == 200  # HTTP 200 Ok
+
+    def test_upload_api_404(self, logged_in_client):
+        # Act
+        response = logged_in_client.post("/api/upload/")
+
+        # Assert
+        assert response.status_code == 400  # HTTP 400 Bad Request
+
+
 @pytest.mark.django_db
 class TestMappingViewSet:
     def test_get_mappings(self, logged_in_client):