CintraAI/code-chunker

1import unittest2from test_base import BaseChunkerTest, mock_count_tokens3from Chunker import CodeChunker4 5class TestCodeChunkerPython(BaseChunkerTest):6    def setUp(self):7        super().setUp()8        self.code_chunker = CodeChunker(file_extension='py')9    10    def test_chunk_simple_code(self):11        py_code = self.mock_codebase['simple.py']12        first_chunk_token_limit = mock_count_tokens("import sys")13        print(f"first_chunk_token_limit = {first_chunk_token_limit}")14        chunks = self.code_chunker.chunk(py_code, token_limit=25)15        token_count = self.mock_count_tokens(py_code)16        print(f"token_count = {token_count}")17        print(f"original code:\n {py_code}")18        19        chunks, full_code = self.run_chunker_test(py_code, token_limit=25)20        21        self.assertEqual(len(chunks), 2) # There should be 2 chunks22        self.assertIn("import sys", chunks[1]) # The first chunk should contain the import statement23        self.assertIn("print('Hello, world!')", chunks[2]) # The second chunk should contain the print statement24 25    def test_chunk_code_text_only(self):26        py_code = self.mock_codebase['text_only.py']27        chunks, final_code = self.run_chunker_test(py_code)28        29        self.assertEqual(len(chunks), 1)30        self.assertIn("This file is empty and should test the chunker's ability to handle empty files", chunks[1])31 32    def test_chunk_code_with_routes(self):33        py_code = self.mock_codebase['routes.py']34        self.run_chunker_test(py_code)35 36    def test_chunk_code_with_models(self):37        py_code = self.mock_codebase['models.py']38        self.run_chunker_test(py_code)39 40    def test_chunk_code_with_main(self):41        py_code = self.mock_codebase['main.py']42        self.run_chunker_test(py_code)43 44    def test_chunk_code_with_utilities(self):45        py_code = self.mock_codebase['utilities.py']46        self.run_chunker_test(py_code)47 48    def test_chunk_code_with_big_class(self):49        py_code = self.mock_codebase['big_class.py']50        self.run_chunker_test(py_code)51 52if __name__ == '__main__':53    unittest.main()