update tests

tests/model_tests.py

@@ -1,11 +1,10 @@
 import unittest
 
-from src.main import app
+from src.main import classify_arabic_dialect
 
 
 class TestClassifier(unittest.TestCase):
     def setUp(self) -> None:
-        self.client = app.test_client()
         self.dialects = ['Egypt', 'Iraq', 'Saudi_Arabia', 'Mauritania', 'Algeria', 'Syria',
         'Oman', 'Tunisia', 'Lebanon', 'Morocco', 'Djibouti','United_Arab_Emirates','Kuwait', 
         'Libya', 'Bahrain', 'Qatar', 'Yemen', 'Palestine', 'Jordan', 'Somalia', 'Sudan']
@@ -17,19 +16,16 @@ class TestClassifier(unittest.TestCase):
         }
 
     def test_response(self):
-        """Test if the response of the /classify API endpoint is correct"""
-        request_data = {"text": "حاجة حلوة اكيد"}
-        response = self.client.post("/classify", json=request_data)
-        self.assertEqual(response.status_code, 200)
-        self.assertIn("class", response.json)
-        self.assertIn(response.json["class"], self.dialects)
+        """Test if the response of the main function is correct"""
+        text = "حاجة حلوة اكيد"
+        response = classify_arabic_dialect(text)
+        self.assertIn(response, self.dialects)
 
     def test_model_output(self):
         """Test that the model correctly classifies obvious dialects"""
         for country, text, in self.test_set.items():
-            request_data = {"text": text}
-            response = self.client.post("/classify", json=request_data)
-            self.assertEqual(response.json["class"], country)
+            response = classify_arabic_dialect(text)
+            self.assertEqual(response, country)
 
     
 if __name__ == "__main__":