ReMe/tests/unit/test_jsonl_file_chunker.py

"""Tests for JsonlFileChunker."""

# pylint: disable=protected-access

import asyncio
import json
import os
import tempfile

from reme.components.file_chunker import JsonlFileChunker

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _run(coro):
    return asyncio.run(coro)


def _write_jsonl(lines: list[str]) -> str:
    """Write lines to a temp .jsonl file and return its path."""
    fd, path = tempfile.mkstemp(suffix=".jsonl")
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        for line in lines:
            f.write(line if line.endswith("\n") else line + "\n")
    return path


def _make_records(n: int, width: int = 20) -> list[str]:
    """Generate *n* JSONL lines, each roughly *width* chars of JSON."""
    return [json.dumps({"id": i, "text": "x" * max(0, width - 20)}) for i in range(n)]


# ---------------------------------------------------------------------------
# Basic behaviour
# ---------------------------------------------------------------------------


def test_empty_file():
    """Empty file → zero chunks."""
    path = _write_jsonl([])
    try:
        chunker = JsonlFileChunker()
        node, chunks = _run(chunker.chunk(path))
        assert len(chunks) == 0
        assert node.links == []
        print("✓ test_empty_file passed")
    finally:
        os.unlink(path)


def test_blank_lines_only():
    """File with only blank lines → zero chunks."""
    fd, path = tempfile.mkstemp(suffix=".jsonl")
    with os.fdopen(fd, "w") as f:
        f.write("\n\n\n")
    try:
        chunker = JsonlFileChunker()
        _, chunks = _run(chunker.chunk(path))
        assert len(chunks) == 0
        print("✓ test_blank_lines_only passed")
    finally:
        os.unlink(path)


def test_single_line():
    """One line → one chunk."""
    lines = _make_records(1)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker()
        _, chunks = _run(chunker.chunk(path))
        assert len(chunks) == 1
        assert chunks[0].start_line == 1
        assert chunks[0].end_line == 1
        assert lines[0] in chunks[0].text
        print("✓ test_single_line passed")
    finally:
        os.unlink(path)


def test_unicode_line_separator_inside_json_string_is_not_a_record_boundary():
    """U+2028 is valid JSON string content, not a JSONL physical newline."""
    record = json.dumps({"id": 1, "text": "before\u2028\u2028after"}, ensure_ascii=False)
    path = _write_jsonl([record])
    try:
        chunker = JsonlFileChunker(max_lines_per_chunk=1)
        node, chunks = _run(chunker.chunk(path))

        assert len(chunks) == 1
        assert chunks[0].text == record + "\n"
        assert chunks[0].start_line == chunks[0].end_line == 1
        assert node.chunk_ids == [chunks[0].id]
    finally:
        os.unlink(path)


def test_all_lines_fit_in_one_chunk():
    """Small file fits entirely in one chunk."""
    lines = _make_records(5, width=30)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=5000)
        _, chunks = _run(chunker.chunk(path))
        assert len(chunks) == 1
        assert chunks[0].start_line == 1
        assert chunks[0].end_line == 5
        print("✓ test_all_lines_fit_in_one_chunk passed")
    finally:
        os.unlink(path)


def test_max_lines_per_chunk_one():
    """max_lines_per_chunk=1 emits exactly one complete line per chunk."""
    lines = _make_records(5, width=10)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=5000, max_lines_per_chunk=1)
        _, chunks = _run(chunker.chunk(path))
        assert len(chunks) == len(lines)
        for line_number, (line, chunk) in enumerate(zip(lines, chunks), start=1):
            assert chunk.start_line == line_number
            assert chunk.end_line == line_number
            assert chunk.text == line + "\n"
    finally:
        os.unlink(path)


# ---------------------------------------------------------------------------
# Multi-chunk splitting (no overlap)
# ---------------------------------------------------------------------------


def test_multiple_chunks_no_overlap():
    """Large file is split into multiple chunks with no overlap."""
    # Each line ≈ 32 chars (with \n), 10 lines ≈ 320 chars.
    lines = _make_records(10, width=30)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=100, max_overlap_chars=0)
        _, chunks = _run(chunker.chunk(path))
        assert len(chunks) > 1, f"Expected >1 chunks, got {len(chunks)}"
        # No overlap: consecutive chunks don't share lines.
        for i in range(len(chunks) - 1):
            assert chunks[i].end_line < chunks[i + 1].start_line
        # All lines are covered.
        assert chunks[0].start_line == 1
        assert chunks[-1].end_line == 10
        print(f"  Created {len(chunks)} chunks")
        print("✓ test_multiple_chunks_no_overlap passed")
    finally:
        os.unlink(path)


def test_line_aligned_no_intra_line_split():
    """Every chunk boundary falls on a line boundary."""
    lines = _make_records(20, width=40)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=150, max_overlap_chars=0)
        _, chunks = _run(chunker.chunk(path))
        assert len(chunks) > 1
        for c in chunks:
            # Each chunk text should be a concatenation of whole lines.
            text_lines = c.text.splitlines()
            assert len(text_lines) == (c.end_line - c.start_line + 1)
        print("✓ test_line_aligned_no_intra_line_split passed")
    finally:
        os.unlink(path)


# ---------------------------------------------------------------------------
# Overlap behaviour
# ---------------------------------------------------------------------------


def test_overlap_lines_shared():
    """Overlap lines appear at the end of chunk N and the start of chunk N+1."""
    # Each line is ~32 chars.  max_chars=100 → ~3 lines per chunk.
    # max_overlap_chars=40 → ~1 line of overlap.
    lines = _make_records(12, width=30)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=100, max_overlap_chars=40)
        _, chunks = _run(chunker.chunk(path))
        assert len(chunks) > 1
        # Verify overlap: chunk[i].end_line >= chunk[i+1].start_line
        for i in range(len(chunks) - 1):
            assert chunks[i].end_line >= chunks[i + 1].start_line, (
                f"Chunk {i} end_line={chunks[i].end_line} < " f"chunk {i+1} start_line={chunks[i+1].start_line}"
            )
        # Last chunk still reaches the end of the file.
        assert chunks[-1].end_line == 12
        print(f"  Created {len(chunks)} chunks with overlap")
        print("✓ test_overlap_lines_shared passed")
    finally:
        os.unlink(path)


def test_zero_overlap():
    """max_overlap_chars=0 produces strictly non-overlapping chunks."""
    lines = _make_records(10, width=30)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=100, max_overlap_chars=0)
        _, chunks = _run(chunker.chunk(path))
        for i in range(len(chunks) - 1):
            assert chunks[i].end_line < chunks[i + 1].start_line
        print("✓ test_zero_overlap passed")
    finally:
        os.unlink(path)


def test_overlap_capped_by_max_overlap_chars():
    """Overlap never exceeds max_overlap_chars in size."""
    lines = _make_records(10, width=50)
    path = _write_jsonl(lines)
    try:
        max_overlap = 60  # Should allow at most 1 line of overlap (~51 chars).
        chunker = JsonlFileChunker(max_chars=200, max_overlap_chars=max_overlap)
        _, chunks = _run(chunker.chunk(path))
        if len(chunks) > 1:
            for i in range(len(chunks) - 1):
                overlap_start = chunks[i + 1].start_line
                overlap_end = chunks[i].end_line
                if overlap_start <= overlap_end:
                    overlap_lines = list(range(overlap_start, overlap_end + 1))
                    overlap_text = "".join(lines[ln - 1] for ln in overlap_lines)
                    assert len(overlap_text) <= max_overlap, f"Overlap size {len(overlap_text)} > {max_overlap}"
        print("✓ test_overlap_capped_by_max_overlap_chars passed")
    finally:
        os.unlink(path)


# ---------------------------------------------------------------------------
# Oversized single line
# ---------------------------------------------------------------------------


def test_single_line_exceeds_max():
    """A line larger than max_chars is emitted alone (no intra-line split)."""
    long_line = json.dumps({"data": "A" * 200})
    lines = ["short", long_line, "short"]
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=64, max_overlap_chars=0)
        _, chunks = _run(chunker.chunk(path))
        # The long line must appear in exactly one chunk by itself (or with
        # no split).
        found = False
        for c in chunks:
            if long_line in c.text:
                found = True
                break
        assert found, "Long line was not found in any chunk"
        assert len(chunks) >= 2  # At least some splitting happened.
        print("✓ test_single_line_exceeds_max passed")
    finally:
        os.unlink(path)


# ---------------------------------------------------------------------------
# Byte mode
# ---------------------------------------------------------------------------


def test_byte_mode():
    """mode='bytes' counts byte length instead of char length."""
    # CJK characters are 3 bytes each in UTF-8 but 1 char each.
    # Use ensure_ascii=False so raw CJK chars appear in the file.
    cjk_record = json.dumps({"text": "你" * 10}, ensure_ascii=False)
    # The JSON string itself has CJK chars; after writing to file each line
    # has ~13 ASCII chars ({"text": "..."}) + 10 CJK chars.
    # char length ≈ 23, byte length ≈ 13 + 10*3 = 43.
    lines = [cjk_record] * 6
    path = _write_jsonl(lines)
    try:
        # In char mode: each line ≈ 24 chars → 2 lines per chunk at max=50.
        chunker_chars = JsonlFileChunker(max_chars=50, max_overlap_chars=0, mode="chars")
        _, chunks_chars = _run(chunker_chars.chunk(path))

        # In byte mode: each line ≈ 44 bytes → 1 line per chunk at max=50.
        chunker_bytes = JsonlFileChunker(max_chars=50, max_overlap_chars=0, mode="bytes")
        _, chunks_bytes = _run(chunker_bytes.chunk(path))

        # Byte mode should produce more chunks since each line is larger in bytes.
        assert len(chunks_bytes) > len(chunks_chars), (
            f"Byte mode ({len(chunks_bytes)}) should produce more chunks "
            f"than char mode ({len(chunks_chars)}) for CJK text"
        )
        print(f"  char mode: {len(chunks_chars)} chunks, byte mode: {len(chunks_bytes)} chunks")
        print("✓ test_byte_mode passed")
    finally:
        os.unlink(path)


def test_invalid_mode_raises():
    """An invalid mode string raises ValueError."""
    try:
        JsonlFileChunker(mode="words")
        assert False, "Should have raised ValueError"
    except ValueError:
        pass
    print("✓ test_invalid_mode_raises passed")


# ---------------------------------------------------------------------------
# FileNode / FileChunk properties
# ---------------------------------------------------------------------------


def test_file_chunk_properties():
    """Each FileChunk has correct path, line range, text, and hash id."""
    lines = _make_records(5, width=30)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=5000)
        _, chunks = _run(chunker.chunk(path))
        c = chunks[0]
        assert c.start_line >= 1
        assert c.end_line >= c.start_line
        assert c.id and len(c.id) > 0
        assert c.text.strip()
        print("✓ test_file_chunk_properties passed")
    finally:
        os.unlink(path)


def test_file_node_properties():
    """FileNode carries path, st_mtime, chunk_ids, and empty links."""
    lines = _make_records(3)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker()
        node, chunks = _run(chunker.chunk(path))
        assert node.st_mtime > 0
        assert node.links == []
        assert node.chunk_ids == [c.id for c in chunks]
        print("✓ test_file_node_properties passed")
    finally:
        os.unlink(path)


def test_hash_id_deterministic():
    """Same file produces the same chunk hash ids on repeated calls."""
    lines = _make_records(5, width=30)
    path = _write_jsonl(lines)
    try:
        c1 = JsonlFileChunker(max_chars=100)
        _, chunks1 = _run(c1.chunk(path))
        c2 = JsonlFileChunker(max_chars=100)
        _, chunks2 = _run(c2.chunk(path))
        assert [c.id for c in chunks1] == [c.id for c in chunks2]
        print("✓ test_hash_id_deterministic passed")
    finally:
        os.unlink(path)


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------


def test_max_chars_floor():
    """max_chars is clamped to at least 64."""
    c = JsonlFileChunker(max_chars=1)
    assert c.max_chars == 64
    print("✓ test_max_chars_floor passed")


def test_max_overlap_floor():
    """Negative max_overlap_chars is clamped to 0."""
    c = JsonlFileChunker(max_overlap_chars=-10)
    assert c.max_overlap_chars == 0
    print("✓ test_max_overlap_floor passed")


def test_full_coverage():
    """Every line of the file appears in at least one chunk."""
    lines = _make_records(25, width=40)
    path = _write_jsonl(lines)
    try:
        chunker = JsonlFileChunker(max_chars=150, max_overlap_chars=50)
        _, chunks = _run(chunker.chunk(path))
        covered: set[int] = set()
        for c in chunks:
            covered.update(range(c.start_line, c.end_line + 1))
        assert covered == set(range(1, 26)), f"Missing lines: {set(range(1, 26)) - covered}"
        print("✓ test_full_coverage passed")
    finally:
        os.unlink(path)


# ---------------------------------------------------------------------------
# Internal algorithm
# ---------------------------------------------------------------------------


def test_chunk_lines_algorithm_simple():
    """_chunk_lines produces correct ranges for a simple case."""
    # max_chars floor is 64, so lines must be long enough.
    # Each line = 20 chars + \n = 21 chars.  64 / 21 = 3 lines per chunk.
    chunker = JsonlFileChunker(max_chars=64, max_overlap_chars=0)
    lines = ["a" * 20 + "\n", "b" * 20 + "\n", "c" * 20 + "\n", "d" * 20 + "\n"]
    # 3 lines = 63 chars ≤ 64 → chunk 1: (0,3).  4th line alone → chunk 2: (3,4).
    ranges = chunker._chunk_lines(lines)
    assert ranges == [(0, 3), (3, 4)], f"Got {ranges}"
    print("✓ test_chunk_lines_algorithm_simple passed")


def test_chunk_lines_algorithm_with_overlap():
    """_chunk_lines with overlap produces overlapping ranges."""
    # Each line = 20 chars + \n = 21 chars.  max_chars=64 → 3 lines per chunk.
    # max_overlap_chars=22 → 1 line of overlap (21 chars ≤ 22).
    chunker = JsonlFileChunker(max_chars=64, max_overlap_chars=22)
    lines = ["a" * 20 + "\n"] * 6
    ranges = chunker._chunk_lines(lines)
    assert len(ranges) >= 2, f"Expected ≥2 ranges, got {ranges}"
    # Verify overlap: range[i].end >= range[i+1].start
    for i in range(len(ranges) - 1):
        assert ranges[i][1] >= ranges[i + 1][0], f"No overlap between range {i} {ranges[i]} and {i+1} {ranges[i+1]}"
    # Last range reaches end.
    assert ranges[-1][1] == 6
    print("✓ test_chunk_lines_algorithm_with_overlap passed")


# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------


if __name__ == "__main__":
    test_empty_file()
    test_blank_lines_only()
    test_single_line()
    test_unicode_line_separator_inside_json_string_is_not_a_record_boundary()
    test_all_lines_fit_in_one_chunk()
    test_multiple_chunks_no_overlap()
    test_line_aligned_no_intra_line_split()
    test_overlap_lines_shared()
    test_zero_overlap()
    test_overlap_capped_by_max_overlap_chars()
    test_single_line_exceeds_max()
    test_byte_mode()
    test_invalid_mode_raises()
    test_file_chunk_properties()
    test_file_node_properties()
    test_hash_id_deterministic()
    test_max_chars_floor()
    test_max_overlap_floor()
    test_full_coverage()
    test_chunk_lines_algorithm_simple()
    test_chunk_lines_algorithm_with_overlap()
    print("\n所有测试通过!")