import os import unittest from io import BytesIO from unittest import mock import app.image_classifier_app as ui_app from PIL import Image path_tests = os.path.dirname(os.path.abspath(__file__)) class TestMLService(unittest.TestCase): # 💡 NOTE Run tests with: python tests/test_image_classifier_app.py def setUp(self): self.token = "dummy_token" self.image_file = Image.open(path_tests + "/dog.jpeg") self.uploaded_file = mock.MagicMock(spec=BytesIO) self.uploaded_file.getvalue.return_value = BytesIO() self.uploaded_file.name = "dog.jpeg" self.image_file.save(self.uploaded_file, format="JPEG") self.headers = {"Authorization": f"Bearer {self.token}"} # python3 -m unittest -vvv tests.test_model def test_login_success(self): # 💡 NOTE Run test with: python -m unittest -vvv tests.test_image_classifier_app.TestMLService.test_login_success expected_token = "dummy_token" response_data = {"access_token": expected_token} with mock.patch("requests.post") as mock_post: mock_post.return_value.status_code = 200 mock_post.return_value.json.return_value = response_data token = ui_app.login("username", "password") headers = { "accept": "application/json", "Content-Type": "application/x-www-form-urlencoded", } self.assertEqual(token, expected_token) mock_post.assert_called_once_with( ui_app.API_BASE_URL + "/login", headers=headers, data={ "grant_type": "", "username": "username", "password": "password", "scope": "", "client_id": "", "client_secret": "", }, ) def test_login_failure(self): # 💡 NOTE Run test with: python -m unittest -vvv tests.test_image_classifier_app.TestMLService.test_login_failure with mock.patch("requests.post") as mock_post: mock_post.return_value.status_code = 401 token = ui_app.login("username", "password") headers = { "accept": "application/json", "Content-Type": "application/x-www-form-urlencoded", } self.assertIsNone(token) mock_post.assert_called_once_with( ui_app.API_BASE_URL + "/login", headers=headers, data={ "grant_type": "", "username": "username", "password": "password", "scope": "", "client_id": "", "client_secret": "", }, ) def test_predict_success(self): # 💡 NOTE Run test with: python -m unittest -vvv tests.test_image_classifier_app.TestMLService.test_predict_success expected_response = {"prediction": "Eskimo_dog", "score": 0.9346} with mock.patch("requests.post") as mock_post: mock_post.return_value.status_code = 200 mock_post.return_value.json.return_value = expected_response response = ui_app.predict(self.token, self.uploaded_file) self.assertEqual(response.status_code, 200) self.assertEqual(response.json(), expected_response) mock_post.assert_called_once_with( ui_app.API_BASE_URL + "/model/predict", files={ "file": (self.uploaded_file.name, self.uploaded_file.getvalue()) }, headers=self.headers, ) def test_predict_failure(self): # 💡 NOTE Run test with: python -m unittest -vvv tests.test_image_classifier_app.TestMLService.test_predict_failure with mock.patch("requests.post") as mock_post: mock_post.return_value.status_code = 500 response = ui_app.predict(self.token, self.uploaded_file) self.assertEqual(response.status_code, 500) mock_post.assert_called_once_with( ui_app.API_BASE_URL + "/model/predict", files={ "file": (self.uploaded_file.name, self.uploaded_file.getvalue()) }, headers=self.headers, ) def test_send_feedback_success(self): # 💡 NOTE Run test with: python -m unittest -vvv tests.test_image_classifier_app.TestMLService.test_send_feedback_success expected_response = {"status": "success"} feedback = "This is a feedback" score = 0.9346 prediction = "Eskimo_dog" image_file_name = "dog.jpeg" with mock.patch("requests.post") as mock_post: mock_post.return_value.status_code = 201 mock_post.return_value.json.return_value = expected_response response = ui_app.send_feedback( self.token, feedback, score, prediction, image_file_name ) self.assertEqual(response.status_code, 201) self.assertEqual(response.json(), expected_response) mock_post.assert_called_once_with( ui_app.API_BASE_URL + "/feedback", json={ "feedback": feedback, "score": score, "predicted_class": prediction, "image_file_name": image_file_name, }, headers=self.headers, ) def test_send_feedback_failure(self): # 💡 NOTE Run test with: python -m unittest -vvv tests.test_image_classifier_app.TestMLService.test_send_feedback_failure feedback = "This is a feedback" score = 0.9346 prediction = "Eskimo_dog" image_file_name = "dog.jpeg" with mock.patch("requests.post") as mock_post: mock_post.return_value.status_code = 500 response = ui_app.send_feedback( self.token, feedback, score, prediction, image_file_name ) self.assertEqual(response.status_code, 500) mock_post.assert_called_once_with( ui_app.API_BASE_URL + "/feedback", json={ "feedback": feedback, "score": score, "predicted_class": prediction, "image_file_name": image_file_name, }, headers=self.headers, ) if __name__ == "__main__": unittest.main(verbosity=2)