Merge pull request #403 from platiagro/fix/handle_exception_prediction

add exception handling

Author: dfvneto

projects/controllers/predictions.py

@@ -5,9 +5,10 @@
 from typing import Optional
 
 import requests
+from sqlalchemy.exc import DataError
 from platiagro import load_dataset
 
-from projects import models, schemas
+from projects import models, schemas, exceptions
 from projects.controllers.utils import (
     parse_dataframe_to_seldon_request,
     parse_file_buffer_to_seldon_request,
@@ -121,7 +122,10 @@ def create_prediction_database_object(
         )
 
         self.session.add(prediction)
-        self.session.commit()
+        try:
+            self.session.commit()
+        except DataError:
+            raise(exceptions.BadRequest("400", "File too large"))
         return prediction
 
     def start_and_save_seldon_prediction(self, request_body, prediction_object, url):

projects/controllers/tasks/tasks.py

@@ -12,7 +12,7 @@
 from jinja2 import Template
 from sqlalchemy import asc, desc, func
 
-from projects import __version__, models, schemas
+from projects import models, schemas
 from projects.controllers.utils import uuid_alpha
 from projects.exceptions import BadRequest, Forbidden, NotFound
 from projects.kubernetes.notebook import (

projects/schemas/project.py

@@ -37,6 +37,7 @@ def validate_description(cls, v):
 class ProjectUpdate(ProjectBase):
     name: Optional[str]
     description: Optional[str]
+
     @validator("name")
     def validate_name(cls, v):
         generic_validators.raise_if_exceeded(MAX_CHARS_ALLOWED, v)

tests/test_experiments.py

@@ -517,7 +517,7 @@ def test_update_experiment_invalid_name(self):
         )
         rv.json()
         self.assertEqual(rv.status_code, 400)
-    
+
     def test_update_experiment_invalid_name_size(self):
         """
         Should return http status 400.
@@ -531,4 +531,4 @@ def test_update_experiment_invalid_name_size(self):
             json={"name": experiment_name},
         )
         rv.json()
-        self.assertEqual(rv.status_code, 400)
+        self.assertEqual(rv.status_code, 400)

tests/test_predictions.py

@@ -3,6 +3,7 @@
 import unittest.mock as mock
 
 from fastapi.testclient import TestClient
+from sqlalchemy.exc import DataError
 
 from projects.api.main import app
 from projects.database import session_scope
@@ -228,3 +229,33 @@ def test_create_prediction_dataset_image(
                 }
             },
         )
+
+    @mock.patch(
+        "projects.controllers.predictions.load_dataset",
+        return_value=util.IRIS_DATAFRAME,
+    )
+    @mock.patch(
+        "requests.post",
+        return_value=util.MOCK_POST_PREDICTION,
+    )
+    @mock.patch("projects.api.predictions.Session.commit", side_effect=DataError("statement", "params", "orig"))
+    def test_create_prediction_fail(
+        self,
+        mock_requests_post,
+        mock_load_dataset,
+        mock_data_error
+    ):
+        """
+        Should return 400 error because the used dataset file is too large.
+        """
+        project_id = util.MOCK_UUID_1
+        deployment_id = util.MOCK_UUID_1
+        dataset_name = util.IRIS_DATASET_NAME
+
+        rv = TEST_CLIENT.post(
+            f"/projects/{project_id}/deployments/{deployment_id}/predictions",
+            json={"dataset": dataset_name},
+        )
+        result = rv.json()
+        self.assertEqual(rv.status_code, 400)
+        self.assertEqual(result["message"], "File too large")

tests/test_projects.py

@@ -20,6 +20,7 @@
                 LoremipsumdolorsitametconsecteturadipiscingelitInteerelitexauc\
                 LoremipsumdolorsitametconsecteturadipiscingelitInteerelitexauc"
 
+
 class TestProjects(unittest.TestCase):
     maxDiff = None
 
@@ -444,7 +445,7 @@ def test_update_project_invalid_name(self):
         rv.json()
 
         self.assertEqual(rv.status_code, 400)
-    
+
     def test_update_project_invalid_name_size(self):
         """
         Should return http status 400.
@@ -465,7 +466,7 @@ def test_update_project_invalid_description_size(self):
         rv.json()
 
         self.assertEqual(rv.status_code, 400)
-    
+
     def test_update_project_description_success(self):
         """
         Should return http status 200.
@@ -474,4 +475,4 @@ def test_update_project_description_success(self):
         rv = TEST_CLIENT.patch(f"/projects/{project_id}", json={"description": "DESCRIPTION"})
         rv.json()
 
-        self.assertEqual(rv.status_code, 200)
+        self.assertEqual(rv.status_code, 200)