Pytest for post /predict endpoint

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+#Pycharm
+.idea

requirements.txt

@@ -8,3 +8,5 @@ pytest==5.4.3
 scikit-learn==0.24
 joblib==0.16.0
 gunicorn==20.0.4
+pytest==6.2.2
+locust==1.4.3

tests/test_response.py

@@ -0,0 +1,13 @@
+from fastapi.testclient import TestClient
+from app import app
+
+
+def test_success_prediction():
+    endpoint = '/predict'
+    body = {"data": [[4.8, 3, 1.4, 0.3], [2, 1, 3.2, 1.1]]}
+
+    with TestClient(app) as client:
+        response = client.post(endpoint, json=body)
+        response_json = response.json()
+        assert response.status_code == 200
+        assert 'prediction' in response_json