Akshanshsensei/PDF-Constrained-Conversational-Agent

1import unittest2from unittest.mock import patch, MagicMock3from core.query_router import QueryRouter4 5class TestQueryRouter(unittest.TestCase):6    def setUp(self):7        self.router = QueryRouter()8 9    def test_page_count_regex(self):10        res = self.router.classify("how many pages are in this document?")11        self.assertEqual(res["intent"], "factual")12        self.assertEqual(res["sub_intent"], "page_count")13        self.assertEqual(res["confidence"], 1.0)14 15    def test_word_count_regex(self):16        res = self.router.classify("how many times is the word 'climate' used?")17        self.assertEqual(res["intent"], "factual")18        self.assertEqual(res["sub_intent"], "word_count:climate")19        self.assertEqual(res["confidence"], 1.0)20 21    def test_aggregation_regex(self):22        res = self.router.classify("can you list all instances of renewable energy?")23        self.assertEqual(res["intent"], "aggregation")24        self.assertIsNone(res["sub_intent"])25        self.assertEqual(res["confidence"], 1.0)26 27    def test_top_words_regex(self):28        res = self.router.classify("what are the top 5 most used words?")29        self.assertEqual(res["intent"], "factual")30        self.assertEqual(res["sub_intent"], "top_words:5")31        self.assertEqual(res["confidence"], 1.0)32 33    @patch('core.query_router.genai.Client')34    def test_explanation_semantic_llm(self, mock_client):35        # Mock LLM response36        mock_response = MagicMock()37        mock_response.text = '{"intent": "semantic", "reasoning": "Asking for explanation"}'38        mock_models = MagicMock()39        mock_models.generate_content.return_value = mock_response40        mock_instance = MagicMock()41        mock_instance.models = mock_models42        mock_client.return_value = mock_instance43 44        res = self.router.classify("Explain the greenhouse effect based on the text.")45        self.assertEqual(res["intent"], "semantic")46        self.assertEqual(res["confidence"], 0.5)47 48    @patch('core.query_router.genai.Client')49    def test_ambiguous_query_default(self, mock_client):50        # Mock LLM returning invalid format or unsupported intent51        mock_response = MagicMock()52        mock_response.text = '{"intent": "unknown"}'53        mock_models = MagicMock()54        mock_models.generate_content.return_value = mock_response55        mock_instance = MagicMock()56        mock_instance.models = mock_models57        mock_client.return_value = mock_instance58 59        res = self.router.classify("what?")60        self.assertEqual(res["intent"], "semantic")61        self.assertEqual(res["confidence"], 0.5)62 63    @patch('core.query_router.genai.Client')64    def test_hindi_query_semantic(self, mock_client):65        mock_response = MagicMock()66        mock_response.text = '{"intent": "semantic", "reasoning": "Hindi explanation"}'67        mock_models = MagicMock()68        mock_models.generate_content.return_value = mock_response69        mock_instance = MagicMock()70        mock_instance.models = mock_models71        mock_client.return_value = mock_instance72 73        res = self.router.classify("इस दस्तावेज़ में क्या है?")74        self.assertEqual(res["intent"], "semantic")75 76    @patch('core.query_router.genai.Client')77    def test_french_query_semantic(self, mock_client):78        mock_response = MagicMock()79        mock_response.text = '{"intent": "semantic", "reasoning": "French explanation"}'80        mock_models = MagicMock()81        mock_models.generate_content.return_value = mock_response82        mock_instance = MagicMock()83        mock_instance.models = mock_models84        mock_client.return_value = mock_instance85 86        res = self.router.classify("Quel est le sujet principal ?")87        self.assertEqual(res["intent"], "semantic")88 89if __name__ == '__main__':90    unittest.main()91