Spaces:
Running
Running
File size: 2,114 Bytes
a50881d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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() |