Spaces:
Running
Running
| import unittest | |
| from test_base import BaseChunkerTest, mock_count_tokens | |
| from Chunker import CodeChunker | |
| class TestCodeChunkerPython(BaseChunkerTest): | |
| def setUp(self): | |
| super().setUp() | |
| self.code_chunker = CodeChunker(file_extension='py') | |
| def test_chunk_simple_code(self): | |
| py_code = self.mock_codebase['simple.py'] | |
| first_chunk_token_limit = mock_count_tokens("import sys") | |
| print(f"first_chunk_token_limit = {first_chunk_token_limit}") | |
| chunks = self.code_chunker.chunk(py_code, token_limit=25) | |
| token_count = self.mock_count_tokens(py_code) | |
| print(f"token_count = {token_count}") | |
| print(f"original code:\n {py_code}") | |
| chunks, full_code = self.run_chunker_test(py_code, token_limit=25) | |
| self.assertEqual(len(chunks), 2) # There should be 2 chunks | |
| self.assertIn("import sys", chunks[1]) # The first chunk should contain the import statement | |
| self.assertIn("print('Hello, world!')", chunks[2]) # The second chunk should contain the print statement | |
| def test_chunk_code_text_only(self): | |
| py_code = self.mock_codebase['text_only.py'] | |
| chunks, final_code = self.run_chunker_test(py_code) | |
| self.assertEqual(len(chunks), 1) | |
| self.assertIn("This file is empty and should test the chunker's ability to handle empty files", chunks[1]) | |
| def test_chunk_code_with_routes(self): | |
| py_code = self.mock_codebase['routes.py'] | |
| self.run_chunker_test(py_code) | |
| def test_chunk_code_with_models(self): | |
| py_code = self.mock_codebase['models.py'] | |
| self.run_chunker_test(py_code) | |
| def test_chunk_code_with_main(self): | |
| py_code = self.mock_codebase['main.py'] | |
| self.run_chunker_test(py_code) | |
| def test_chunk_code_with_utilities(self): | |
| py_code = self.mock_codebase['utilities.py'] | |
| self.run_chunker_test(py_code) | |
| def test_chunk_code_with_big_class(self): | |
| py_code = self.mock_codebase['big_class.py'] | |
| self.run_chunker_test(py_code) | |
| if __name__ == '__main__': | |
| unittest.main() |