arborist/tests/test_compress.py

"""Tests for transparent zstd compression of chunk content."""

from __future__ import annotations

import sqlite3

from aborist.compress import (
    _MIN_COMPRESS_BYTES,
    _ZSTD_MAGIC,
    is_compressed,
    pack_chunk,
    unpack_chunk,
)
from aborist.ingest import ingest_source
from aborist.sources.wikipedia_xml import WikipediaXmlDump
from aborist.store import connect


def test_pack_passes_short_text_through_uncompressed():
    short = "hello world"
    packed = pack_chunk(short)
    assert packed == short
    assert isinstance(packed, str)
    assert not is_compressed(packed)


def test_pack_compresses_long_text():
    long = "Wikipedia article body. " * 200  # ~5 KB, well above threshold
    packed = pack_chunk(long)
    assert isinstance(packed, bytes)
    assert is_compressed(packed)
    assert len(packed) < len(long.encode("utf-8"))


def test_pack_falls_back_to_plaintext_when_incompressible():
    # Already-random bytes hex'd as text — no entropy left for zstd.
    import os
    incompressible = os.urandom(2048).hex()
    packed = pack_chunk(incompressible)
    # Implementation may keep plaintext if compression doesn't shrink.
    if isinstance(packed, bytes):
        assert len(packed) < len(incompressible.encode("utf-8"))
    else:
        assert packed == incompressible


def test_unpack_round_trip_compressed():
    text = "Some readable text here. " * 100
    packed = pack_chunk(text)
    assert is_compressed(packed)
    assert unpack_chunk(packed) == text


def test_unpack_passes_legacy_str_through():
    # Old DBs have plaintext str in chunks.content. Must keep working.
    legacy = "old-format plaintext row"
    assert unpack_chunk(legacy) == legacy


def test_unpack_handles_none_for_cold_tier():
    assert unpack_chunk(None) is None


def test_unpack_handles_plain_utf8_bytes_without_magic():
    raw = "no zstd here, just bytes".encode("utf-8")
    assert not is_compressed(raw)
    assert unpack_chunk(raw) == "no zstd here, just bytes"


def test_unpack_raises_on_unknown_type():
    import pytest

    with pytest.raises(TypeError):
        unpack_chunk(12345)


def test_zstd_magic_is_correct_rfc8478_value():
    # Sanity: zstd frame magic per RFC 8478 §3.1.1.
    assert _ZSTD_MAGIC == b"\x28\xb5\x2f\xfd"


def test_min_compress_threshold_is_reasonable():
    # Threshold below which the helper passes plaintext through.
    # 64 bytes is plenty small; covers TF-IDF cores (~50-200 bytes) but
    # would trigger compression for anything article-shaped.
    assert _MIN_COMPRESS_BYTES <= 256


def test_ingest_writes_compressed_for_large_chunks(tmp_path):
    """End-to-end: ingest a doc with a long body, verify chunks.content
    is stored as compressed bytes (not plaintext)."""
    # Synthesize a fixture XML with one big article body.
    big_body = " ".join(f"sentence{i}" for i in range(2000))  # ~18 KB
    xml = f"""<?xml version="1.0" encoding="UTF-8"?>
<mediawiki xmlns="http://www.mediawiki.org/xml/export-0.10/" version="0.10">
  <page>
    <title>Big Article</title>
    <ns>0</ns>
    <id>1</id>
    <revision>
      <id>1</id>
      <timestamp>2010-01-01T00:00:00Z</timestamp>
      <text xml:space="preserve">{big_body}</text>
    </revision>
  </page>
</mediawiki>"""
    fixture = tmp_path / "wp.xml"
    fixture.write_text(xml, encoding="utf-8")

    db_path = tmp_path / "aborist.db"
    conn = connect(db_path)
    try:
        ingest_source(conn, WikipediaXmlDump(fixture))
    finally:
        conn.close()

    # Re-open with raw sqlite3 to inspect cell types directly.
    raw = sqlite3.connect(db_path)
    row = raw.execute("SELECT content FROM chunks LIMIT 1").fetchone()
    raw.close()
    assert row is not None
    content = row[0]
    # Stored as bytes (BLOB cell), zstd-framed.
    assert isinstance(content, bytes)
    assert is_compressed(content)
    # Round-trips back to the original text.
    assert unpack_chunk(content).startswith("sentence0")


def test_ingest_then_search_round_trip_finds_compressed_doc(tmp_path):
    """FTS5 still indexes plaintext; chunks.content stays compressed; reads
    that go through unpack_chunk reconstruct the original. The whole
    pipeline keeps working."""
    from aborist.search import FTS5Backend

    big_body = "merkle providence wikipedia anarchism " * 200  # ~7.4 KB
    xml = f"""<?xml version="1.0" encoding="UTF-8"?>
<mediawiki xmlns="http://www.mediawiki.org/xml/export-0.10/" version="0.10">
  <page>
    <title>Topic</title>
    <ns>0</ns>
    <id>1</id>
    <revision>
      <id>1</id>
      <timestamp>2010-01-01T00:00:00Z</timestamp>
      <text xml:space="preserve">{big_body}</text>
    </revision>
  </page>
</mediawiki>"""
    fixture = tmp_path / "wp.xml"
    fixture.write_text(xml, encoding="utf-8")

    db_path = tmp_path / "aborist.db"
    conn = connect(db_path)
    try:
        ingest_source(conn, WikipediaXmlDump(fixture))
        backend = FTS5Backend(conn)
        hits = backend.search("merkle providence")
    finally:
        conn.close()
    assert hits, "FTS5 should find the compressed doc via its plaintext index"
    assert hits[0].title == "Topic"