|
| 1 | +# Copyright (c) 2024 Microsoft Corporation. |
| 2 | +# Licensed under the MIT License |
| 3 | + |
| 4 | +from unittest.mock import Mock, patch |
| 5 | + |
| 6 | +from graphrag.config.models.chunking_config import ChunkingConfig |
| 7 | +from graphrag.index.bootstrap import bootstrap |
| 8 | +from graphrag.index.operations.chunk_text.strategies import run_sentences, run_tokens |
| 9 | +from graphrag.index.operations.chunk_text.typing import TextChunk |
| 10 | + |
| 11 | + |
| 12 | +class TestRunSentences: |
| 13 | + def setup_method(self, method): |
| 14 | + bootstrap() |
| 15 | + |
| 16 | + def test_basic_functionality(self): |
| 17 | + """Test basic sentence splitting without metadata""" |
| 18 | + input = ["This is a test. Another sentence."] |
| 19 | + tick = Mock() |
| 20 | + chunks = list(run_sentences(input, ChunkingConfig(), tick)) |
| 21 | + |
| 22 | + assert len(chunks) == 2 |
| 23 | + assert chunks[0].text_chunk == "This is a test." |
| 24 | + assert chunks[1].text_chunk == "Another sentence." |
| 25 | + assert all(c.source_doc_indices == [0] for c in chunks) |
| 26 | + tick.assert_called_once_with(1) |
| 27 | + |
| 28 | + def test_multiple_documents(self): |
| 29 | + """Test processing multiple input documents""" |
| 30 | + input = ["First. Document.", "Second. Doc."] |
| 31 | + tick = Mock() |
| 32 | + chunks = list(run_sentences(input, ChunkingConfig(), tick)) |
| 33 | + |
| 34 | + assert len(chunks) == 4 |
| 35 | + assert chunks[0].source_doc_indices == [0] |
| 36 | + assert chunks[2].source_doc_indices == [1] |
| 37 | + assert tick.call_count == 2 |
| 38 | + |
| 39 | + def test_mixed_whitespace_handling(self): |
| 40 | + """Test input with irregular whitespace""" |
| 41 | + input = [" Sentence with spaces. Another one! "] |
| 42 | + chunks = list(run_sentences(input, ChunkingConfig(), Mock())) |
| 43 | + assert chunks[0].text_chunk == " Sentence with spaces." |
| 44 | + assert chunks[1].text_chunk == "Another one!" |
| 45 | + |
| 46 | + |
| 47 | +class TestRunTokens: |
| 48 | + @patch("tiktoken.get_encoding") |
| 49 | + def test_basic_functionality(self, mock_get_encoding): |
| 50 | + mock_encoder = Mock() |
| 51 | + mock_encoder.encode.side_effect = lambda x: list(x.encode()) |
| 52 | + mock_encoder.decode.side_effect = lambda x: bytes(x).decode() |
| 53 | + mock_get_encoding.return_value = mock_encoder |
| 54 | + |
| 55 | + # Input and config |
| 56 | + input = [ |
| 57 | + "Marley was dead: to begin with. There is no doubt whatever about that. The register of his burial was signed by the clergyman, the clerk, the undertaker, and the chief mourner. Scrooge signed it. And Scrooge's name was good upon 'Change, for anything he chose to put his hand to." |
| 58 | + ] |
| 59 | + config = ChunkingConfig(size=5, overlap=1, encoding_model="fake-encoding") |
| 60 | + tick = Mock() |
| 61 | + |
| 62 | + # Run the function |
| 63 | + chunks = list(run_tokens(input, config, tick)) |
| 64 | + |
| 65 | + # Verify output |
| 66 | + assert len(chunks) > 0 |
| 67 | + assert all(isinstance(chunk, TextChunk) for chunk in chunks) |
| 68 | + tick.assert_called_once_with(1) |
| 69 | + |
| 70 | + @patch("tiktoken.get_encoding") |
| 71 | + def test_non_string_input(self, mock_get_encoding): |
| 72 | + """Test handling of non-string input (e.g., numbers).""" |
| 73 | + mock_encoder = Mock() |
| 74 | + mock_encoder.encode.side_effect = lambda x: list(str(x).encode()) |
| 75 | + mock_encoder.decode.side_effect = lambda x: bytes(x).decode() |
| 76 | + mock_get_encoding.return_value = mock_encoder |
| 77 | + |
| 78 | + input = [123] # Non-string input |
| 79 | + config = ChunkingConfig(size=5, overlap=1, encoding_model="fake-encoding") |
| 80 | + tick = Mock() |
| 81 | + |
| 82 | + chunks = list(run_tokens(input, config, tick)) # type: ignore |
| 83 | + |
| 84 | + # Verify non-string input is handled |
| 85 | + assert len(chunks) > 0 |
| 86 | + assert "123" in chunks[0].text_chunk |
0 commit comments