Akshanshsensei/PDF-Constrained-Conversational-Agent

1import unittest2from core.tools import _count_word, _get_page_count, _get_page_content, _find_all_occurrences, ToolExecutor3from core.metadata_store import DocumentMetadata4 5class TestTools(unittest.TestCase):6    def setUp(self):7        self.full_text = "Hello world. This is a test document. The word 'test' appears twice. No wait, testing doesn't count. Test!"8        self.page_boundaries = [0, 50, 100] # Let's say 3 pages9 10    def test_count_word(self):11        # Case insensitive exact match.12        # "test" -> "test document", "word 'test'", "Test!" -> 3 matches13        # "testing" does not count.14        count = _count_word("test", self.full_text)15        self.assertEqual(count, 3)16        17        # Word not in text18        self.assertEqual(_count_word("banana", self.full_text), 0)19 20    def test_get_page_count(self):21        self.assertEqual(_get_page_count(3), 3)22 23    def test_get_page_content(self):24        # bounds: 0, 50, 10025        # len is 10826        page_1 = _get_page_content(1, self.full_text, self.page_boundaries)27        page_2 = _get_page_content(2, self.full_text, self.page_boundaries)28        page_3 = _get_page_content(3, self.full_text, self.page_boundaries)29        30        self.assertEqual(page_1, self.full_text[0:50].strip())31        self.assertEqual(page_2, self.full_text[50:100].strip())32        self.assertEqual(page_3, self.full_text[100:108].strip())33        34        # Out of bounds35        with self.assertRaises(ValueError):36            _get_page_content(0, self.full_text, self.page_boundaries)37        with self.assertRaises(ValueError):38            _get_page_content(4, self.full_text, self.page_boundaries)39 40    def test_find_all_occurrences(self):41        results = _find_all_occurrences("test", self.full_text, self.page_boundaries)42        self.assertEqual(len(results), 4) # "test document", "test' appears", "testing doesn't", "Test!"43        44        # Check snippet truncation and page numbers45        # "test document" is around index 23 -> page 146        self.assertEqual(results[0]["page"], 1)47        self.assertTrue("test" in results[0]["snippet"].lower())48 49    def test_tool_executor(self):50        metadata = DocumentMetadata(3, 10, 108, {}, self.page_boundaries, self.full_text)51        executor = ToolExecutor(metadata)52        53        result = executor.execute("count_word", word="test")54        self.assertEqual(result.output, 3)55        self.assertIsNone(result.error)56        57        # Unknown tool58        result2 = executor.execute("fake_tool")59        self.assertIsNotNone(result2.error)60