CintraAI/code-chunker

1import unittest2from unittest.mock import patch3import tiktoken4import sys5import os6 7# Add parent directory to path to allow imports8sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))9 10from Chunker import Chunker, CodeChunker11from utils import load_json12 13# Mocking the count_tokens function as it's external and not the focus of these tests14def mock_count_tokens(string: str, encoding_name='gpt-4') -> int:15    """Returns the number of tokens in a text string."""16    encoding = tiktoken.encoding_for_model(encoding_name)17    num_tokens = len(encoding.encode(string))18    return num_tokens19 20class BaseChunkerTest(unittest.TestCase):21    """Base class for all code chunker tests with common setup and utilities."""22    23    def setUp(self):24        self.patcher = patch('utils.count_tokens', side_effect=mock_count_tokens)25        self.mock_count_tokens = self.patcher.start()26        self.mock_codebase = load_json('mock_codefiles.json')27        28    def tearDown(self):29        self.patcher.stop()30    31    def run_chunker_test(self, code, token_limit=20):32        """Helper method to run standard chunker tests."""33        chunks = self.code_chunker.chunk(code, token_limit=token_limit)34        Chunker.print_chunks(chunks)35        final_code = Chunker.consolidate_chunks_into_file(chunks)36        num_lines = Chunker.count_lines(final_code)37        38        # Common assertions39        self.assertEqual(num_lines, len(code.split("\n")))40        self.assertIn(code, final_code)41        42        return chunks, final_code