Alex Galarce
674add9a3e
Feature/black and isort * black and isort constants * black and isort codes * black and isort crosswalks * black and isort investment * black and isort prompts * isort testbed * isort tracking * black and isort utils * poetry and black tests Approved-by: Katon Minhas
73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
import pytest
|
|
from src.investment.smart_chunking_funcs import parse_chunk, stitch_chunks
|
|
|
|
|
|
def test_parse_chunk_valid():
|
|
chunk = "chunk 001: This is a test chunk."
|
|
chunk_number, content = parse_chunk(chunk)
|
|
assert chunk_number == 1
|
|
assert content == "This is a test chunk."
|
|
|
|
|
|
def test_parse_chunk_valid_with_large_number():
|
|
chunk = "chunk 123: Another test chunk with a larger number."
|
|
chunk_number, content = parse_chunk(chunk)
|
|
assert chunk_number == 123
|
|
assert content == "Another test chunk with a larger number."
|
|
|
|
|
|
def test_parse_chunk_empty_content():
|
|
chunk = "chunk 005: "
|
|
chunk_number, content = parse_chunk(chunk)
|
|
assert chunk_number == 5
|
|
assert content == ""
|
|
|
|
|
|
def test_parse_chunk_invalid_format():
|
|
chunk = "invalid chunk format"
|
|
with pytest.raises(ValueError):
|
|
parse_chunk(chunk)
|
|
|
|
|
|
def test_parse_chunk_short_prefix():
|
|
chunk = "chunk 01: Short prefix"
|
|
with pytest.raises(ValueError):
|
|
parse_chunk(chunk)
|
|
|
|
|
|
def test_stitch_chunks_empty_list():
|
|
assert stitch_chunks([]) == ""
|
|
|
|
|
|
def test_stitch_chunks_single_chunk():
|
|
assert stitch_chunks(["chunk 001: hello world"]) == "hello world"
|
|
|
|
|
|
def test_stitch_chunks_two_chunks():
|
|
chunks = ["chunk 001: hello world", "chunk 002: world and more"]
|
|
assert stitch_chunks(chunks, overlap_size=5) == "hello world and more"
|
|
|
|
|
|
def test_stitch_chunks_out_of_order():
|
|
chunks = ["chunk 002: world", "chunk 001: hello"]
|
|
assert stitch_chunks(chunks, overlap_size=0) in ["helloworld", "hello\nworld"]
|
|
|
|
|
|
def test_stitch_noncontiguous_chunks():
|
|
chunks = ["chunk 001: hello", "chunk 003: world"]
|
|
assert stitch_chunks(chunks, overlap_size=10) in [
|
|
"helloworld",
|
|
"hello\nworld",
|
|
] # the overlap size will be ignored in the noncontiguous case
|
|
|
|
|
|
def test_stitch_long_overlap(): # not enough text to overlap, so we just append them
|
|
chunks = ["chunk 001: hello", "chunk 002: world"]
|
|
assert stitch_chunks(chunks, overlap_size=10) in ["helloworld", "hello\nworld"]
|
|
|
|
|
|
def test_stitch_chunks_negative_overlap():
|
|
chunks = ["chunk 001: hello", "chunk 002: world"]
|
|
with pytest.raises(ValueError):
|
|
stitch_chunks(chunks, overlap_size=-1)
|