doczyai-pipelines/fieldExtraction/tests/test_smart_chunking_funcs.py

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)