Add test case

TogetherCrew · Sep 13, 2024 · 18ce27e · 18ce27e
1 parent 34ea417
commit 18ce27e
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 4 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1 +1,2 @@
-__pycache__
+__pycache__
+.pytest_cache
diff --git a/main.py b/main.py
@@ -6,8 +6,8 @@
 pipe = pipeline("text-classification", model="shahrukhx01/question-vs-statement-classifier")
 
 custom_labels = {
-    "LABEL_0": "statement",
-    "LABEL_1": "question"
+    "LABEL_0": "STATEMENT",
+    "LABEL_1": "QUESTION"
 }
 
 class Payload(BaseModel):

diff --git a/requirements.txt b/requirements.txt
@@ -1,3 +1,5 @@
 fastapi[standard]==0.114.1
 pydantic==2.8.0
-transformers[torch]==4.44.2
+transformers[torch]==4.44.2
+pytest==8.3.3
+httpx==0.27.2
diff --git a/test_main.py b/test_main.py
@@ -0,0 +1,10 @@
+from fastapi.testclient import TestClient
+from main import app
+
+client = TestClient(app)
+
+def test_route():
+  response = client.post('/test', json={ "text": "Is this a test?" })
+  assert response.status_code == 200
+  assert response.json()['label'] == 'QUESTION'
+