Katon Minhas
c4e519894b
Refactor/one time io * Delete client work * Streamline imports * Fix list class * Update tests * Clear code funcs test * Fix unit tests * Single read code funcs * Single-read model * Single-load embeddings * Successful E2E test * update unit tests * Update importsg * Reload poetry.lock * remove old exhibit header function * Remove aarete_derived generic function * remove align and format tables * Remove strings to dict * references * Clear code * Move preprocessing_funcs * remove keywords * refactor postprocessing_funcs * Pass unit test - remove qa_qc directory * Black and isort * remove print Approved-by: Alex Galarce
74 lines
2.2 KiB
Python
74 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)
|