Add exception tests

Author: Alex-Welsh

web-apps/chat/test.py

@@ -3,9 +3,9 @@
 import unittest
 
 from gradio_client import Client
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch, MagicMock, Mock
 from langchain.schema import HumanMessage, AIMessage, SystemMessage
-from app import build_chat_context, inference, PossibleSystemPromptException, BACKEND_INITIALISED
+from app import build_chat_context, inference, PossibleSystemPromptException, gr
 
 url = os.environ.get("GRADIO_URL", "http://localhost:7860")
 client = Client(url)
@@ -95,6 +95,111 @@ def test_inference_thinking_tags(self, mock_build_chat_context, mock_llm):
 
         self.assertEqual(responses, ["Thinking...", "Thinking...", "", "final response"])
 
+    @patch("app.llm")
+    @patch("app.INCLUDE_SYSTEM_PROMPT", True)
+    @patch("app.build_chat_context")
+    def test_inference_PossibleSystemPromptException(self, mock_build_chat_context, mock_llm):
+        mock_build_chat_context.return_value = ["mock_context"]
+        mock_response = Mock()
+        mock_response.json.return_value = {"message": "Bad request"}
+
+        mock_llm.stream.side_effect = openai.BadRequestError(
+            message="Bad request",
+            response=mock_response,
+            body=None
+        )
+
+        latest_message = "Hello"
+        history = []
+
+        with self.assertRaises(PossibleSystemPromptException):
+            list(inference(latest_message, history))
+
+    @patch("app.llm")
+    @patch("app.INCLUDE_SYSTEM_PROMPT", False)
+    @patch("app.build_chat_context")
+    def test_inference_general_error(self, mock_build_chat_context, mock_llm):
+        mock_build_chat_context.return_value = ["mock_context"]
+        mock_response = Mock()
+        mock_response.json.return_value = {"message": "Bad request"}
+
+        mock_llm.stream.side_effect = openai.BadRequestError(
+            message="Bad request",
+            response=mock_response,
+            body=None
+        )
+
+        latest_message = "Hello"
+        history = []
+        exception_message = "\'API Error received. This usually means the chosen LLM uses an incompatible prompt format. Error message was: Bad request\'"
+
+        with self.assertRaises(gr.Error) as gradio_error:
+            list(inference(latest_message, history))
+        self.assertEqual(str(gradio_error.exception), exception_message)
+
+    @patch("app.llm")
+    @patch("app.build_chat_context")
+    @patch("app.log")
+    @patch("app.gr")
+    @patch("app.BACKEND_INITIALISED", False)
+    def test_inference_APIConnectionError(self, mock_gr, mock_logger, mock_build_chat_context, mock_llm):
+        mock_build_chat_context.return_value = ["mock_context"]
+        mock_request = Mock()
+        mock_request.json.return_value = {"message": "Foo"}
+
+        mock_llm.stream.side_effect = openai.APIConnectionError(
+            message="Foo",
+            request=mock_request,
+        )
+
+        latest_message = "Hello"
+        history = []
+
+        list(inference(latest_message, history))
+        mock_logger.info.assert_any_call("Backend API not yet ready")
+        mock_gr.Info.assert_any_call("Backend not ready - model may still be initialising - please try again later.")
+
+    @patch("app.llm")
+    @patch("app.build_chat_context")
+    @patch("app.log")
+    @patch("app.gr")
+    @patch("app.BACKEND_INITIALISED", True)
+    def test_inference_APIConnectionError_initialised(self, mock_gr, mock_logger, mock_build_chat_context, mock_llm):
+        mock_build_chat_context.return_value = ["mock_context"]
+        mock_request = Mock()
+        mock_request.json.return_value = {"message": "Foo"}
+
+        mock_llm.stream.side_effect = openai.APIConnectionError(
+            message="Foo",
+            request=mock_request,
+        )
+
+        latest_message = "Hello"
+        history = []
+
+        list(inference(latest_message, history))
+        mock_logger.error.assert_called_once_with("Failed to connect to backend API: %s", mock_llm.stream.side_effect)
+        mock_gr.Warning.assert_any_call("Failed to connect to backend API.")
+
+    @patch("app.llm")
+    @patch("app.build_chat_context")
+    @patch("app.gr")
+    def test_inference_InternalServerError(self, mock_gr, mock_build_chat_context, mock_llm):
+        mock_build_chat_context.return_value = ["mock_context"]
+        mock_request = Mock()
+        mock_request.json.return_value = {"message": "Foo"}
+
+        mock_llm.stream.side_effect = openai.InternalServerError(
+            message="Foo",
+            response=mock_request,
+            body=None
+        )
+
+        latest_message = "Hello"
+        history = []
+
+        list(inference(latest_message, history))
+        mock_gr.Warning.assert_any_call("Internal server error encountered in backend API - see API logs for details.")
 
 if __name__ == "__main__":
     unittest.main()