Various testing tweaks

Added client_message and history as global vars to reduce duplication
Increased verbosity
Created separate classes for different test categories

Author: Alex-Welsh

web-apps/chat/test.py

@@ -9,69 +9,68 @@
 
 url = os.environ.get("GRADIO_URL", "http://localhost:7860")
 client = Client(url)
-
-class TestSuite(unittest.TestCase):
-    # General tests
+latest_message = "Why don't humans drink horse milk?"
+history = [
+    {
+        "role": "user",
+        "metadata": None,
+        "content": "Hi!",
+        "options": None,
+    },
+    {
+        "role": "assistant",
+        "metadata": None,
+        "content": "Hello! How can I help you?",
+        "options": None,
+    },
+]
+
+class TestAPI(unittest.TestCase):
     def test_gradio_api(self):
         result = client.predict("Hi", api_name="/chat")
         self.assertGreater(len(result), 0)
 
-    # build_chat_context function tests
+class TestBuildChatContext(unittest.TestCase):
     @patch("app.settings")
     @patch("app.INCLUDE_SYSTEM_PROMPT", True)
     def test_chat_context_system_prompt(self, mock_settings):
         mock_settings.model_instruction = "You are a helpful assistant."
-        latest_message = "What is a mammal?"
-        history = [
-            {'role': 'user', 'metadata': None, 'content': 'Hi!', 'options': None},
-            {"role": "assistant", 'metadata': None, "content": "Hello! How can I help you?", 'options': None},
-        ]
 
         context = build_chat_context(latest_message, history)
 
         self.assertEqual(len(context), 4)
         self.assertIsInstance(context[0], SystemMessage)
         self.assertEqual(context[0].content, "You are a helpful assistant.")
         self.assertIsInstance(context[1], HumanMessage)
-        self.assertEqual(context[1].content, "Hi!")
+        self.assertEqual(context[1].content, history[0]["content"])
         self.assertIsInstance(context[2], AIMessage)
-        self.assertEqual(context[2].content, "Hello! How can I help you?")
+        self.assertEqual(context[2].content, history[1]["content"])
         self.assertIsInstance(context[3], HumanMessage)
         self.assertEqual(context[3].content, latest_message)
 
     @patch("app.settings")
     @patch("app.INCLUDE_SYSTEM_PROMPT", False)
     def test_chat_context_human_prompt(self, mock_settings):
         mock_settings.model_instruction = "You are a very helpful assistant."
-        latest_message = "What is a fish?"
-        history = [
-            {"role": "user", 'metadata': None, "content": "Hi there!", 'options': None},
-            {"role": "assistant", 'metadata': None, "content": "Hi! How can I help you?", 'options': None},
-        ]
 
         context = build_chat_context(latest_message, history)
 
         self.assertEqual(len(context), 3)
         self.assertIsInstance(context[0], HumanMessage)
-        self.assertEqual(context[0].content, "You are a very helpful assistant.\n\nHi there!")
+        self.assertEqual(context[0].content, "You are a very helpful assistant.\n\nHi!")
         self.assertIsInstance(context[1], AIMessage)
-        self.assertEqual(context[1].content, "Hi! How can I help you?")
+        self.assertEqual(context[1].content, history[1]["content"])
         self.assertIsInstance(context[2], HumanMessage)
         self.assertEqual(context[2].content, latest_message)
 
-    # inference function tests
+class TestInference(unittest.TestCase):
     @patch("app.settings")
     @patch("app.llm")
     @patch("app.log")
     def test_inference_success(self, mock_logger, mock_llm, mock_settings):
         mock_llm.stream.return_value = [MagicMock(content="response_chunk")]
 
         mock_settings.model_instruction = "You are a very helpful assistant."
-        latest_message = "Why don't we drink horse milk?"
-        history = [
-            {"role": "user", 'metadata': None, "content": "Hi there!", 'options': None},
-            {"role": "assistant", 'metadata': None, "content": "Hi! How can I help you?", 'options': None},
-        ]
 
         responses = list(inference(latest_message, history))
 
@@ -88,8 +87,6 @@ def test_inference_thinking_tags(self, mock_build_chat_context, mock_llm):
             MagicMock(content="</think>"),
             MagicMock(content="final response"),
         ]
-        latest_message = "Hello"
-        history = []
 
         responses = list(inference(latest_message, history))
 
@@ -98,7 +95,8 @@ def test_inference_thinking_tags(self, mock_build_chat_context, mock_llm):
     @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):
+    @patch("app.log")
+    def test_inference_PossibleSystemPromptException(self, mock_logger, 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"}
@@ -109,16 +107,15 @@ def test_inference_PossibleSystemPromptException(self, mock_build_chat_context,
             body=None
         )
 
-        latest_message = "Hello"
-        history = []
-
         with self.assertRaises(PossibleSystemPromptException):
             list(inference(latest_message, history))
+        mock_logger.error.assert_called_once_with("Received BadRequestError from backend API: %s", mock_llm.stream.side_effect)
 
     @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):
+    @patch("app.log")
+    def test_inference_general_error(self, mock_logger, 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"}
@@ -129,13 +126,12 @@ def test_inference_general_error(self, mock_build_chat_context, mock_llm):
             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)
+        mock_logger.error.assert_called_once_with("Received BadRequestError from backend API: %s", mock_llm.stream.side_effect)
 
     @patch("app.llm")
     @patch("app.build_chat_context")
@@ -152,9 +148,6 @@ def test_inference_APIConnectionError(self, mock_gr, mock_logger, mock_build_cha
             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.")
@@ -174,9 +167,6 @@ def test_inference_APIConnectionError_initialised(self, mock_gr, mock_logger, mo
             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.")
@@ -195,11 +185,8 @@ def test_inference_InternalServerError(self, mock_gr, mock_build_chat_context, m
             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()
+    unittest.main(verbosity=2)

web-apps/test-images.sh

@@ -76,7 +76,7 @@ test() {
             --name $1-test-suite \
             -e GRADIO_URL=http://$1-app:7860 --entrypoint python \
             $IMAGE \
-            test.py
+            test.py -v
 
         log "Removing containers:"
         docker rm -f ollama $1-app