Updates

main.py

@@ -44,8 +44,8 @@ async def classify_image(file: UploadFile = File(...)):
 
         # Prepare the custom response data
         response_data = {
-            "filename": file.filename,
-            "isNsfw": best_prediction['label'] == 'nsfw',
+            "file_name": file.filename,
+            "is_nsfw": best_prediction['label'] == 'nsfw',
             "confidence_percentage": confidence_percentage
         }
 

test-requirements.txt

@@ -0,0 +1,3 @@
+fastapi==0.110.2
+httpx==0.27.0
+pytest==8.1.1

test_main.py

@@ -0,0 +1,16 @@
+from fastapi.testclient import TestClient
+from main import app
+
+client = TestClient(app)
+
+file_name = "sunflower.jpg"
+
+def test_read_main():
+    response = client.post("/api/v1/detect", files={"file": (file_name, open(file_name, "rb"), "image/jpeg")})
+    assert response.status_code == 200
+    assert response.json() == {"file_name": file_name, "is_nsfw": False, "confidence_percentage": 100.0}
+
+def test_invalid_input():
+    response = client.post("/api/v1/detect", files={})
+    assert response.status_code == 422 
+