NexusOS/tests/test_documents.py

"""Document ingest / RAG store — hermetic (fake embeddings, temp DB)."""
import asyncio
import tempfile
from pathlib import Path

from synapse.memory.store import PersistentMemoryStore


def _store():
    return PersistentMemoryStore(Path(tempfile.mkdtemp()) / "t.db")


def test_chunker_packs_and_splits():
    s = _store()
    one = s._chunk_text("short one.\n\nshort two.")
    assert one == ["short one.\n\nshort two."]  # both fit one chunk
    many = s._chunk_text("a" * 700 + "\n\n" + "b" * 700)
    assert len(many) == 2  # each paragraph near the size cap -> own chunk


async def _fake_embed(text):
    kws = ["lego", "star", "wars", "gpu", "vega"]
    v = [float(text.lower().count(k)) for k in kws]
    return v if any(v) else None


def test_add_list_search_delete_roundtrip():
    async def run():
        s = _store()
        r = await s.add_document(
            "Guide",
            "Beat the lego star wars boss with the force.\n\nUnrelated gpu vega notes.",
            _fake_embed,
        )
        assert r["chunks"] >= 1
        assert [d["title"] for d in s.list_documents()] == ["Guide"]

        hits = await s.search_documents("lego star wars", _fake_embed, limit=2, min_score=0.1)
        assert hits and "lego" in hits[0]["text"].lower()
        # scores are sorted descending
        assert all(hits[i]["score"] >= hits[i + 1]["score"] for i in range(len(hits) - 1))

        # get_document returns ordered chunks for the viewer
        chunks = s.get_document(r["doc_id"])
        assert [c["chunk_idx"] for c in chunks] == list(range(len(chunks)))

        assert s.delete_document(r["doc_id"]) is True
        assert s.list_documents() == []
        assert s.get_document(r["doc_id"]) == []       # gone -> no chunks
        assert s.delete_document(r["doc_id"]) is False  # already gone

    asyncio.run(run())


def test_search_empty_query_returns_nothing():
    s = _store()
    assert asyncio.run(s.search_documents("", _fake_embed)) == []


def test_extract_text_by_type():
    from synapse.main import _extract_text
    # plain text / markdown -> UTF-8 decode
    assert _extract_text("notes.md", b"# Title\n\nbody") == "# Title\n\nbody"
    assert _extract_text("x.txt", "café".encode("utf-8")) == "café"
    # a real (tiny) PDF built with pypdf -> text extracted back out
    from pypdf import PdfWriter, PdfReader
    import io
    w = PdfWriter()
    w.add_blank_page(width=200, height=200)
    buf = io.BytesIO(); w.write(buf)
    out = _extract_text("blank.pdf", buf.getvalue())
    assert isinstance(out, str)  # blank page -> "" or whitespace, never raises
    assert PdfReader(io.BytesIO(buf.getvalue())).pages  # sanity: it was a valid PDF