Spaces:
Build error
Build error
| """ | |
| Unit tests for response_generator module | |
| Tests LLM response generation functionality | |
| """ | |
| import unittest | |
| from unittest.mock import MagicMock, Mock, patch | |
| from src.response_generator import (_generate_response_impl, | |
| format_chat_history, | |
| generate_xeno_response) | |
| class TestResponseGenerator(unittest.TestCase): | |
| """Test cases for response_generator module""" | |
| def setUp(self): | |
| """Set up test fixtures""" | |
| self.context = """Knowledge Entry 1: | |
| Q: How do I create an account? | |
| A: Visit our website and click Sign Up. | |
| ----------------------------------------""" | |
| self.question = "How can I create an account?" | |
| self.chat_history = [ | |
| {"role": "user", "content": "Hello"}, | |
| {"role": "assistant", "content": "Hi! How can I help you?"}, | |
| ] | |
| def test_format_chat_history(self): | |
| """Test formatting chat history""" | |
| formatted = format_chat_history(self.chat_history) | |
| # Check format | |
| self.assertIn("User: Hello", formatted) | |
| self.assertIn("Assistant: Hi! How can I help you?", formatted) | |
| self.assertIn("\n", formatted) | |
| def test_format_chat_history_empty(self): | |
| """Test formatting empty chat history""" | |
| formatted = format_chat_history([]) | |
| self.assertEqual(formatted, "No previous conversation") | |
| def test_format_chat_history_single_message(self): | |
| """Test formatting single message""" | |
| history = [{"role": "user", "content": "Hello"}] | |
| formatted = format_chat_history(history) | |
| self.assertEqual(formatted, "User: Hello") | |
| def test_format_chat_history_missing_fields(self): | |
| """Test formatting with missing fields""" | |
| history = [ | |
| {"role": "user"}, # Missing content | |
| {"content": "Test"}, # Missing role | |
| ] | |
| formatted = format_chat_history(history) | |
| self.assertIn("User:", formatted) | |
| self.assertIn("Unknown:", formatted) | |
| def test_generate_response_impl(self, mock_genai_client): | |
| """Test internal response generation implementation""" | |
| # Configure mock response | |
| mock_genai_client.models.generate_content.return_value.text = "You can create an account by visiting our website." | |
| response = _generate_response_impl( | |
| self.context, self.question, self.chat_history | |
| ) | |
| # Verify generate_content was called with model and content | |
| mock_genai_client.models.generate_content.assert_called_once() | |
| call_kwargs = mock_genai_client.models.generate_content.call_args[1] | |
| self.assertIn("model", call_kwargs) | |
| self.assertIn("contents", call_kwargs) | |
| # Check response | |
| self.assertEqual(response, "You can create an account by visiting our website.") | |
| def test_generate_response_with_empty_history(self, mock_genai_client): | |
| """Test generating response with empty history""" | |
| mock_genai_client.models.generate_content.return_value.text = "Test response" | |
| response = _generate_response_impl(self.context, self.question, []) | |
| # Verify it still works | |
| self.assertEqual(response, "Test response") | |
| # Check that "None" was used for history in prompt | |
| call_kwargs = mock_genai_client.models.generate_content.call_args[1] | |
| prompt = call_kwargs["contents"] | |
| self.assertIn("None", prompt) | |
| def test_prompt_structure(self, mock_genai_client): | |
| """Test that prompt includes all necessary components""" | |
| mock_genai_client.models.generate_content.return_value.text = "Test response" | |
| _generate_response_impl(self.context, self.question, self.chat_history) | |
| # Get the prompt that was sent | |
| call_kwargs = mock_genai_client.models.generate_content.call_args[1] | |
| prompt = call_kwargs["contents"] | |
| # Verify prompt structure | |
| self.assertIn("HISTORY", prompt) | |
| self.assertIn("CONTEXT", prompt) | |
| self.assertIn("QUESTION", prompt) | |
| self.assertIn(self.context, prompt) | |
| self.assertIn(self.question, prompt) | |
| def test_generate_xeno_response_with_timer(self, mock_genai_client): | |
| """Test generate_xeno_response with timer""" | |
| mock_genai_client.models.generate_content.return_value.text = "Test response" | |
| mock_timer = Mock() | |
| mock_timer.time_step = MagicMock() | |
| mock_timer.time_step.return_value.__enter__ = Mock() | |
| mock_timer.time_step.return_value.__exit__ = Mock() | |
| response = generate_xeno_response( | |
| self.context, self.question, self.chat_history, timer=mock_timer | |
| ) | |
| # Verify timer was used | |
| mock_timer.time_step.assert_called_once_with("llm_generation") | |
| # Verify response | |
| self.assertEqual(response, "Test response") | |
| def test_response_text_stripping(self, mock_genai_client): | |
| """Test that response text is stripped of whitespace""" | |
| mock_genai_client.models.generate_content.return_value.text = "Test response with spaces" | |
| response = _generate_response_impl(self.context, self.question, []) | |
| # Response should be returned as-is from mock | |
| self.assertEqual(response, "Test response with spaces") | |
| def test_system_prompt_inclusion(self, mock_genai_client): | |
| """Test that system prompt is included in generated prompt""" | |
| mock_genai_client.models.generate_content.return_value.text = "Test" | |
| _generate_response_impl(self.context, self.question, []) | |
| # Get the prompt | |
| call_kwargs = mock_genai_client.models.generate_content.call_args[1] | |
| prompt = call_kwargs["contents"] | |
| # Should contain system prompt text | |
| self.assertIn("XENO Support Assistant", prompt) | |
| if __name__ == "__main__": | |
| unittest.main() | |