CintraAI/code-chunker

1import unittest2from unittest.mock import patch3from Chunker import Chunker, CodeChunker4import tiktoken5from utils import load_json6 7 8                        9# Mocking the count_tokens function as it's external and not the focus of these tests10def mock_count_tokens(string: str, encoding_name='gpt-4') -> int:11    """Returns the number of tokens in a text string."""12    encoding = tiktoken.encoding_for_model(encoding_name)13    num_tokens = len(encoding.encode(string))14    return num_tokens15 16# Python Test Class17class TestCodeChunkerPython(unittest.TestCase):18    def setUp(self):19        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)20        self.mock_count_tokens = self.patcher.start()21        self.code_chunker = CodeChunker(file_extension='py')22        self.mock_codebase = load_json('mock_codefiles.json')23        24    def tearDown(self):25        self.patcher.stop()26 27    def test_chunk_simple_code(self):28        py_code = self.mock_codebase['simple.py']29        first_chunk_token_limit = mock_count_tokens("import sys")30        print(f"first_chunk_token_limit = {first_chunk_token_limit}")31        chunks = self.code_chunker.chunk(py_code, token_limit=25)32        token_count = self.mock_count_tokens(py_code)33        print(f"token_count = {token_count}")34        print(f"original code:\n {py_code}")35        Chunker.print_chunks(chunks)36        full_code = Chunker.consolidate_chunks_into_file(chunks)37        print(f"code after consolidation:\n {full_code}")38        num_lines = Chunker.count_lines(full_code)39        self.assertEqual(num_lines, len(py_code.split("\n"))) # The number of lines should be the same40        self.assertIn(full_code, py_code) # The full code should be in the original code41        self.assertEqual(len(chunks), 2) # There should be 2 chunks42        self.assertIn("import sys", chunks[1]) # The first chunk should contain the import statement43        self.assertIn("print('Hello, world!')", chunks[2]) # The second chunk should contain the print statement44 45    def test_chunk_code_text_only(self):46        py_code = self.mock_codebase['text_only.py']47        chunks = self.code_chunker.chunk(py_code, token_limit=20)48        Chunker.print_chunks(chunks)49        final_code = Chunker.consolidate_chunks_into_file(chunks)50        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))51        self.assertEqual(num_lines, len(py_code.split("\n"))) # The number of lines should be the same52        self.assertIn(py_code, final_code) # The full code should be in the original code53        self.assertEqual(len(chunks), 1)54        self.assertIn("This file is empty and should test the chunker's ability to handle empty files", chunks[1])55 56 57    def test_chunk_code_with_routes(self):58        py_code = self.mock_codebase['routes.py']59        chunks = self.code_chunker.chunk(py_code, token_limit=20)60        Chunker.print_chunks(chunks)61        final_code = Chunker.consolidate_chunks_into_file(chunks)62        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))63        self.assertEqual(num_lines, len(py_code.split("\n"))) # The number of lines should be the same64        self.assertIn(py_code, final_code) # The full code should be in the original code65 66 67    def test_chunk_code_with_models(self):68        py_code = self.mock_codebase['models.py']69        chunks = self.code_chunker.chunk(py_code, token_limit=20)70        Chunker.print_chunks(chunks)71        final_code = Chunker.consolidate_chunks_into_file(chunks)72        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))73        self.assertEqual(num_lines, len(py_code.split("\n")))74        self.assertIn(py_code, final_code)75 76    def test_chunk_code_with_main(self):77        py_code = self.mock_codebase['main.py']78        chunks = self.code_chunker.chunk(py_code, token_limit=20)79        Chunker.print_chunks(chunks)80        final_code = Chunker.consolidate_chunks_into_file(chunks)81        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))82        self.assertEqual(num_lines, len(py_code.split("\n")))83        self.assertIn(py_code, final_code)84 85    def test_chunk_code_with_utilities(self):86        py_code = self.mock_codebase['utilities.py']87        chunks = self.code_chunker.chunk(py_code, token_limit=20)88        Chunker.print_chunks(chunks)89        final_code = Chunker.consolidate_chunks_into_file(chunks)90        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))91        self.assertEqual(num_lines, len(py_code.split("\n")))92        self.assertIn(py_code, final_code)93 94    def test_chunk_code_with_big_class(self):95        py_code = self.mock_codebase['big_class.py']96        chunks = self.code_chunker.chunk(py_code, token_limit=20)97        Chunker.print_chunks(chunks)98        final_code = Chunker.consolidate_chunks_into_file(chunks)99        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))100        self.assertEqual(num_lines, len(py_code.split("\n")))101        self.assertIn(py_code, final_code)102 103# JavaScript Test Class104class TestCodeChunkerJavaScript(unittest.TestCase):105 106    def setUp(self):107        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)108        self.mock_count_tokens = self.patcher.start()109        self.code_chunker = CodeChunker(file_extension='js')110        self.mock_codebase = load_json('mock_codefiles.json')111 112    def tearDown(self):113        self.patcher.stop()114 115    def test_chunk_javascript_simple_code(self):116        js_code = self.mock_codebase['simple.js']117        chunks = self.code_chunker.chunk(js_code, token_limit=20)118        Chunker.print_chunks(chunks)119        final_code = Chunker.consolidate_chunks_into_file(chunks)120        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))121        self.assertEqual(num_lines, len(js_code.split("\n")))122        self.assertIn(js_code, final_code)123 124 125    def test_chunk_javascript_with_routes(self):126        js_code = self.mock_codebase['routes.js']127        chunks = self.code_chunker.chunk(js_code, token_limit=20)128        Chunker.print_chunks(chunks)129        final_code = Chunker.consolidate_chunks_into_file(chunks)130        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))131        self.assertEqual(num_lines, len(js_code.split("\n")))132        self.assertIn(js_code, final_code)133 134 135    def test_chunk_javascript_with_models(self):136        js_code = self.mock_codebase['models.js']137        chunks = self.code_chunker.chunk(js_code, token_limit=20)138        Chunker.print_chunks(chunks)139        final_code = Chunker.consolidate_chunks_into_file(chunks)140        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))141        self.assertEqual(num_lines, len(js_code.split("\n")))142        self.assertIn(js_code, final_code)143 144    def test_chunk_javascript_with_main(self):145        js_code = self.mock_codebase['main.js']146        chunks = self.code_chunker.chunk(js_code, token_limit=20)147        Chunker.print_chunks(chunks)148        final_code = Chunker.consolidate_chunks_into_file(chunks)149        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))150        self.assertEqual(num_lines, len(js_code.split("\n")))151        self.assertIn(js_code, final_code)152 153    def test_chunk_javascript_with_utilities(self):154        js_code = self.mock_codebase['utilities.js']155        chunks = self.code_chunker.chunk(js_code, token_limit=20)156        Chunker.print_chunks(chunks)157        final_code = Chunker.consolidate_chunks_into_file(chunks)158        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))159        self.assertEqual(num_lines, len(js_code.split("\n")))160        self.assertIn(js_code, final_code)161 162    def test_chunk_javascript_with_big_class(self):163        js_code = self.mock_codebase['big_class.js']164        chunks = self.code_chunker.chunk(js_code, token_limit=20)165        Chunker.print_chunks(chunks)166        final_code = Chunker.consolidate_chunks_into_file(chunks)167        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))168        self.assertEqual(num_lines, len(js_code.split("\n")))169        self.assertIn(js_code, final_code)170 171    def test_chunk_javascript_with_react_component(self):172        js_code = self.mock_codebase['react_component.js']173        chunks = self.code_chunker.chunk(js_code, token_limit=20)174        Chunker.print_chunks(chunks)175        final_code = Chunker.consolidate_chunks_into_file(chunks)176        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))177        self.assertEqual(num_lines, len(js_code.split("\n")))178        self.assertIn(js_code, final_code)179 180# CSS Test Class181class TestCodeChunkerCSS(unittest.TestCase):182   183    def setUp(self):184        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)185        self.mock_count_tokens = self.patcher.start()186        self.code_chunker = CodeChunker(file_extension='css')187         #Load the JSON data188        self.mock_codebase = load_json('mock_codefiles.json')189 190    def tearDown(self):191        self.patcher.stop()192 193    def test_chunk_css_with_media_query(self):194        css_code = self.mock_codebase['media_queries.css']195        chunks = self.code_chunker.chunk(css_code, token_limit=20)196        Chunker.print_chunks(chunks)197        final_code = Chunker.consolidate_chunks_into_file(chunks)198        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))199        self.assertEqual(num_lines, len(css_code.split("\n")))200        self.assertIn(css_code, final_code)201 202    def test_chunk_css_with_simple_css(self):203        css_code = self.mock_codebase['simple_styles.css']204        chunks = self.code_chunker.chunk(css_code, token_limit=20)205        Chunker.print_chunks(chunks)206        final_code = Chunker.consolidate_chunks_into_file(chunks)207        num_lines = Chunker.count_lines(Chunker.consolidate_chunks_into_file(chunks))208        self.assertEqual(num_lines, len(css_code.split("\n")))209        self.assertIn(css_code, final_code)210 211 212    213# TypeScript Test Class214class TestCodeChunkerTypeScript(unittest.TestCase):215 216    def setUp(self):217        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)218        self.mock_count_tokens = self.patcher.start()219        self.code_chunker = CodeChunker(file_extension='ts')220        self.mock_codebase = load_json('mock_codefiles.json')221 222 223    def tearDown(self):224        self.patcher.stop()225 226    def test_chunk_typescript_code(self):227        ts_code = self.mock_codebase['example.ts']228        chunks = self.code_chunker.chunk(ts_code, token_limit=20)229        Chunker.print_chunks(chunks)230        final_code = Chunker.consolidate_chunks_into_file(chunks)231        num_lines = Chunker.count_lines(final_code)232        self.assertEqual(num_lines, len(ts_code.split("\n")))233        self.assertIn(ts_code, final_code)234        self.assertGreater(len(chunks), 1)  # Ensure the code is actually chunked235 236# Ruby Test Class237class TestCodeChunkerRuby(unittest.TestCase):238 239    def setUp(self):240        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)241        self.mock_count_tokens = self.patcher.start()242        self.code_chunker = CodeChunker(file_extension='rb')243        self.mock_codebase = load_json('mock_codefiles.json')244 245 246    def tearDown(self):247        self.patcher.stop()248 249    def test_chunk_ruby_code(self):250        rb_code = self.mock_codebase['example.rb']251        chunks = self.code_chunker.chunk(rb_code, token_limit=20)252        Chunker.print_chunks(chunks)253        final_code = Chunker.consolidate_chunks_into_file(chunks)254        num_lines = Chunker.count_lines(final_code)255        self.assertEqual(num_lines, len(rb_code.split("\n")))256        self.assertIn(rb_code, final_code)257        self.assertGreater(len(chunks), 1)  # Ensure the code is actually chunked258 259# PHP Test Class260class TestCodeChunkerPHP(unittest.TestCase):261 262    def setUp(self):263        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)264        self.mock_count_tokens = self.patcher.start()265        self.code_chunker = CodeChunker(file_extension='php')266        self.mock_codebase = load_json('mock_codefiles.json')267 268    def tearDown(self):269        self.patcher.stop()270 271    def test_chunk_php_code(self):272        php_code = self.mock_codebase['example.php']273        chunks = self.code_chunker.chunk(php_code, token_limit=20)274        Chunker.print_chunks(chunks)275        final_code = Chunker.consolidate_chunks_into_file(chunks)276        num_lines = Chunker.count_lines(final_code)277        self.assertEqual(num_lines, len(php_code.split("\n")))278        self.assertIn(php_code, final_code)279        self.assertGreater(len(chunks), 1)  # Ensure the code is actually chunked280 281# Golang Test Class282class TestCodeChunkerGolang(unittest.TestCase):283 284    def setUp(self):285        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)286        self.mock_count_tokens = self.patcher.start()287        self.code_chunker = CodeChunker(file_extension='go')288        self.mock_codebase = load_json('mock_codefiles.json')289 290    def tearDown(self):291        self.patcher.stop()292 293    def test_chunk_golang_simple_code(self):294        go_code = self.mock_codebase['simple.go']295        chunks = self.code_chunker.chunk(go_code, token_limit=20)296        Chunker.print_chunks(chunks)297        final_code = Chunker.consolidate_chunks_into_file(chunks)298        num_lines = Chunker.count_lines(final_code)299        self.assertEqual(num_lines, len(go_code.split("\n")))300        self.assertIn(go_code, final_code)301        self.assertGreater(len(chunks), 1)  # Ensure the code is actually chunked302 303    def test_chunk_golang_with_structs(self):304        go_code = self.mock_codebase['structs.go']305        chunks = self.code_chunker.chunk(go_code, token_limit=20)306        Chunker.print_chunks(chunks)307        final_code = Chunker.consolidate_chunks_into_file(chunks)308        num_lines = Chunker.count_lines(final_code)309        self.assertEqual(num_lines, len(go_code.split("\n")))310        self.assertIn(go_code, final_code)311        self.assertGreater(len(chunks), 1)312 313    def test_chunk_golang_with_interfaces(self):314        go_code = self.mock_codebase['interfaces.go']315        chunks = self.code_chunker.chunk(go_code, token_limit=20)316        Chunker.print_chunks(chunks)317        final_code = Chunker.consolidate_chunks_into_file(chunks)318        num_lines = Chunker.count_lines(final_code)319        self.assertEqual(num_lines, len(go_code.split("\n")))320        self.assertIn(go_code, final_code)321        self.assertGreater(len(chunks), 1)322 323    def test_chunk_golang_with_goroutines(self):324        go_code = self.mock_codebase['goroutines.go']325        chunks = self.code_chunker.chunk(go_code, token_limit=20)326        Chunker.print_chunks(chunks)327        final_code = Chunker.consolidate_chunks_into_file(chunks)328        num_lines = Chunker.count_lines(final_code)329        self.assertEqual(num_lines, len(go_code.split("\n")))330        self.assertIn(go_code, final_code)331        self.assertGreater(len(chunks), 1)332 333if __name__ == '__main__':334    unittest.main()335