110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from yuxi.channels.sdk.chunking import ChunkConfig, chunk_message
|
||
|
|
|
||
|
|
|
||
|
|
class TestChunkConfig:
|
||
|
|
def test_default_values(self):
|
||
|
|
cfg = ChunkConfig()
|
||
|
|
assert cfg.limit == 4096
|
||
|
|
assert cfg.mode == "text"
|
||
|
|
|
||
|
|
def test_custom_values(self):
|
||
|
|
cfg = ChunkConfig(limit=2000, mode="markdown")
|
||
|
|
assert cfg.limit == 2000
|
||
|
|
assert cfg.mode == "markdown"
|
||
|
|
|
||
|
|
|
||
|
|
class TestChunkMessageTextMode:
|
||
|
|
def test_short_text_no_chunk(self):
|
||
|
|
result = chunk_message("hello world")
|
||
|
|
assert len(result) == 1
|
||
|
|
assert result[0] == "hello world"
|
||
|
|
|
||
|
|
def test_exact_limit(self):
|
||
|
|
text = "a" * 4096
|
||
|
|
cfg = ChunkConfig(limit=4096)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) == 1
|
||
|
|
assert len(result[0]) == 4096
|
||
|
|
|
||
|
|
def test_empty_text(self):
|
||
|
|
result = chunk_message("")
|
||
|
|
assert len(result) == 1
|
||
|
|
assert result[0] == ""
|
||
|
|
|
||
|
|
def test_over_limit_splits_at_paragraph(self):
|
||
|
|
para1 = "A" * 3000
|
||
|
|
para2 = "B" * 3000
|
||
|
|
text = f"{para1}\n\n{para2}"
|
||
|
|
cfg = ChunkConfig(limit=4096)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) == 2
|
||
|
|
|
||
|
|
def test_over_limit_splits_at_sentence(self):
|
||
|
|
text = ("A" * 1000 + "\u3002") * 5
|
||
|
|
cfg = ChunkConfig(limit=4096)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) == 2
|
||
|
|
|
||
|
|
def test_chunk_size_less_than_limit(self):
|
||
|
|
text = "B" * 5000
|
||
|
|
cfg = ChunkConfig(limit=4096)
|
||
|
|
for ch in chunk_message(text, cfg):
|
||
|
|
assert len(ch) <= 4096, f"Chunk length {len(ch)} exceeds limit"
|
||
|
|
|
||
|
|
def test_multiple_paragraphs_merge_when_possible(self):
|
||
|
|
text = "short1\n\nshort2\n\nshort3"
|
||
|
|
result = chunk_message(text)
|
||
|
|
assert len(result) == 1
|
||
|
|
|
||
|
|
def test_very_long_single_paragraph_splits(self):
|
||
|
|
text = "word " * 5000
|
||
|
|
cfg = ChunkConfig(limit=1000)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) > 1
|
||
|
|
for ch in result:
|
||
|
|
assert len(ch) <= 1000
|
||
|
|
|
||
|
|
def test_no_double_newline_splits_by_size(self):
|
||
|
|
text = "A" * 8000
|
||
|
|
cfg = ChunkConfig(limit=4096)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) >= 2
|
||
|
|
for ch in result:
|
||
|
|
assert len(ch) <= 4096
|
||
|
|
|
||
|
|
|
||
|
|
class TestChunkMessageMarkdownMode:
|
||
|
|
def test_short_markdown_no_chunk(self):
|
||
|
|
text = "# Title\n\nSome content"
|
||
|
|
cfg = ChunkConfig(mode="markdown", limit=4096)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) == 1
|
||
|
|
|
||
|
|
def test_code_block_preserved(self):
|
||
|
|
code_block = "```python\n" + "x = 1\n" * 200 + "```"
|
||
|
|
text = f"Intro\n\n{code_block}\n\nOutro"
|
||
|
|
cfg = ChunkConfig(mode="markdown", limit=500)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) >= 1
|
||
|
|
|
||
|
|
def test_splits_between_sections(self):
|
||
|
|
text = "# Section 1\n\n" + "A" * 3000 + "\n\n# Section 2\n\n" + "B" * 3000
|
||
|
|
cfg = ChunkConfig(mode="markdown", limit=4096)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) >= 2
|
||
|
|
|
||
|
|
def test_markdown_empty_text(self):
|
||
|
|
result = chunk_message("", ChunkConfig(mode="markdown"))
|
||
|
|
assert len(result) == 1
|
||
|
|
assert result[0] == ""
|
||
|
|
|
||
|
|
def test_fenced_code_block_toggle(self):
|
||
|
|
text = "before\n```\ncode\n```\nafter"
|
||
|
|
cfg = ChunkConfig(mode="markdown", limit=4096)
|
||
|
|
result = chunk_message(text, cfg)
|
||
|
|
assert len(result) == 1
|
||
|
|
assert "```" in result[0]
|