Newer
Older
navi-1 / tests / unit / export / test_builder.py
"""Export builder: window math, deterministic content id, encryption roundtrip."""

from __future__ import annotations

import gzip
import json
from datetime import datetime, timezone

import pyrage
import pytest

from navi.config import Settings
from navi.export import builder

TS = datetime(2026, 11, 1, 12, 0, tzinfo=timezone.utc)


class _Tx:
    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False


class FakeConn:
    """Just enough of an asyncpg connection for the builder's query sequence."""

    def __init__(self, sessions: list[str], rows: list[dict], total_sessions: int):
        self._sessions = sessions
        self._rows = rows
        self._total = total_sessions
        self.calls: list[tuple[str, tuple]] = []

    def transaction(self, **kwargs):
        return _Tx()

    async def fetch(self, sql, *args):
        self.calls.append((sql, args))
        assert args[0] == ["user", "assistant"]
        return [{"session_id": s} for s in self._sessions]

    async def fetchval(self, sql, *args):
        return self._total

    def cursor(self, sql, *args, prefetch=500):
        self.calls.append((sql, args))
        conn = self

        async def gen():
            for row in conn._rows:
                yield row

        return gen()


class FakePool:
    def __init__(self, conn):
        self._conn = conn

    def acquire(self):
        pool = self

        class _Acq:
            async def __aenter__(self):
                return pool._conn

            async def __aexit__(self, *exc):
                return False

        return _Acq()


def _row(session: str, seq: int, role: str, content: str) -> dict:
    return {
        "session_id": session,
        "sequence_number": seq,
        "role": role,
        "content": content,
        "created_at": TS,
    }


@pytest.fixture
def export_settings(tmp_path, monkeypatch):
    identity = pyrage.x25519.Identity.generate()
    s = Settings(
        _env_file=None,
        navi_persona_file="",
        navi_export_dir=str(tmp_path / "export_files"),
        navi_export_token="t" * 32,
        navi_export_age_recipient=str(identity.to_public()),
        navi_export_roles="user,assistant",
        navi_export_max_bytes=10_000,
    )
    monkeypatch.setattr(builder, "settings", s)
    return s, identity


@pytest.mark.asyncio
async def test_build_writes_encrypted_jsonl_and_manifest(export_settings):
    _settings, identity = export_settings
    rows = [_row("s1", 1, "user", "привет"), _row("s1", 2, "assistant", 'he said "hi"')]
    pool = FakePool(FakeConn(["s1"], rows, total_sessions=1))

    manifest = await builder.build_export(pool)

    age_path, manifest_path = builder.export_paths()
    assert age_path.exists() and manifest_path.exists()
    assert manifest["rows"] == 2
    assert manifest["sessions"] == 1
    assert manifest["window"]["truncated"] is False
    assert manifest["fields"] == builder.FIELDS
    assert manifest["age_recipient"] == str(identity.to_public())
    assert manifest_path.read_text().count("content_id") == 1

    decrypted = pyrage.decrypt(age_path.read_bytes(), [identity])
    lines = [json.loads(line) for line in gzip.decompress(decrypted).decode().splitlines()]
    assert lines == [
        {"v": 1, "session": "s1", "seq": 1, "role": "user", "content": "привет", "ts": TS.isoformat()},
        {"v": 1, "session": "s1", "seq": 2, "role": "assistant", "content": 'he said "hi"', "ts": TS.isoformat()},
    ]
    # The plaintext never stays on disk.
    assert not list(builder.export_paths()[0].parent.glob("*.tmp"))
    assert not list(builder.export_paths()[0].parent.glob("*.plain*"))


@pytest.mark.asyncio
async def test_same_corpus_keeps_the_ciphertext_and_its_content_id(export_settings):
    _, _ = export_settings
    rows = [_row("s1", 1, "user", "hello")]
    first = await builder.build_export(FakePool(FakeConn(["s1"], rows, 1)))
    age_path, _ = builder.export_paths()
    before = age_path.read_bytes()
    mtime = age_path.stat().st_mtime_ns

    second = await builder.build_export(FakePool(FakeConn(["s1"], rows, 1)))

    assert second["content_id"] == first["content_id"]
    assert age_path.read_bytes() == before
    assert age_path.stat().st_mtime_ns == mtime
    # built_at still moves, so freshness monitoring sees a live builder.
    assert second["built_at"] >= first["built_at"]


@pytest.mark.asyncio
async def test_changed_corpus_produces_a_new_content_id(export_settings):
    _, _ = export_settings
    first = await builder.build_export(FakePool(FakeConn(["s1"], [_row("s1", 1, "user", "a")], 1)))
    second = await builder.build_export(FakePool(FakeConn(["s1"], [_row("s1", 1, "user", "b")], 1)))
    assert first["content_id"] != second["content_id"]


@pytest.mark.asyncio
async def test_window_reports_truncation(export_settings):
    _, _ = export_settings
    manifest = await builder.build_export(
        FakePool(FakeConn(["newest"], [_row("newest", 1, "user", "x")], total_sessions=7))
    )
    assert manifest["window"]["truncated"] is True
    assert manifest["sessions"] == 1
    assert manifest["sessions_total"] == 7


@pytest.mark.asyncio
async def test_unconfigured_export_refuses_to_build(tmp_path, monkeypatch):
    s = Settings(_env_file=None, navi_persona_file="", navi_export_age_recipient="", navi_export_token="")
    monkeypatch.setattr(builder, "settings", s)
    with pytest.raises(RuntimeError):
        await builder.build_export(FakePool(FakeConn([], [], 0)))