Newer
Older
navi-1 / tests / unit / export / test_cli.py
"""End-to-end check of the consumer CLI against the wire contract.

A stub server stands in for Navi (same paths, same bearer check, same manifest
shape) and the CLI is driven as a module so exit codes are asserted directly.
This is the contract test for `scripts/navi_export.py`.
"""

from __future__ import annotations

import gzip
import hashlib
import importlib.util
import json
import threading
from contextlib import contextmanager
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import pyrage

SCRIPT = Path(__file__).resolve().parents[3] / "scripts" / "navi_export.py"
TOKEN = "s" * 40
CORPUS = (
    b'{"v":1,"session":"s1","seq":1,"role":"user","content":"hi","ts":null}\n'
    b'{"v":1,"session":"s1","seq":2,"role":"assistant","content":"hello","ts":null}\n'
)


def _cli():
    spec = importlib.util.spec_from_file_location("navi_export_cli", SCRIPT)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _pack(identity):
    cipher = pyrage.encrypt(gzip.compress(CORPUS, mtime=0), [identity.to_public()])
    manifest = {
        "export_version": 1,
        "built_at": "2026-11-04T12:00:00+00:00",
        "content_id": hashlib.sha256(CORPUS).hexdigest(),
        "file": {
            "name": "messages-v1.jsonl.gz.age",
            "bytes": len(cipher),
            "sha256": hashlib.sha256(cipher).hexdigest(),
        },
        "rows": 2,
        "sessions": 1,
        "sessions_total": 1,
        "window": {"from": None, "to": None, "truncated": False, "max_bytes": 1},
        "age_recipient": str(identity.to_public()),
    }
    return manifest, cipher


@contextmanager
def _serve(manifest: dict, cipher: bytes):
    state = {"downloads": 0}

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *args):
            pass

        def _reply(self, code: int, body: bytes = b"", ctype="application/json"):
            self.send_response(code)
            self.send_header("Content-Type", ctype)
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.headers.get("Authorization") != f"Bearer {TOKEN}":
                self._reply(401)
                return
            if self.path == "/export/v1/manifest":
                self._reply(200, json.dumps(manifest).encode())
            elif self.path == "/export/v1/messages":
                state["downloads"] += 1
                self._reply(200, cipher, "application/octet-stream")
            else:
                self._reply(404)

    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        yield f"http://127.0.0.1:{server.server_address[1]}", state
    finally:
        server.shutdown()
        server.server_close()


def _pull(url, identity_path, out, token=TOKEN):
    return _cli().main(
        [
            "pull",
            "--url", url,
            "--token", token,
            "--identity", str(identity_path),
            "--out", str(out),
        ]
    )


def _identity(tmp_path) -> Path:
    path = tmp_path / "key.txt"
    path.write_text(str(pyrage.x25519.Identity.generate()) + "\n")
    return path


def test_pull_downloads_decrypts_and_installs(tmp_path):
    identity = pyrage.x25519.Identity.generate()
    key = tmp_path / "key.txt"
    key.write_text(str(identity) + "\n")
    manifest, cipher = _pack(identity)
    out = tmp_path / "corpus" / "messages.jsonl.gz"

    with _serve(manifest, cipher) as (url, state):
        assert _pull(url, key, out) == 0

    assert state["downloads"] == 1
    assert gzip.decompress(out.read_bytes()) == CORPUS
    assert json.loads((out.parent / (out.name + ".state.json")).read_text())["content_id"] == (
        manifest["content_id"]
    )
    assert not list(out.parent.glob("*.part*"))


def test_second_pull_downloads_nothing(tmp_path):
    identity = pyrage.x25519.Identity.generate()
    key = tmp_path / "key.txt"
    key.write_text(str(identity) + "\n")
    manifest, cipher = _pack(identity)
    out = tmp_path / "messages.jsonl.gz"

    with _serve(manifest, cipher) as (url, state):
        assert _pull(url, key, out) == 0
        assert _pull(url, key, out) == 3
        assert state["downloads"] == 1


def test_wrong_token_exits_1(tmp_path):
    identity = pyrage.x25519.Identity.generate()
    key = tmp_path / "key.txt"
    key.write_text(str(identity) + "\n")
    manifest, cipher = _pack(identity)

    with _serve(manifest, cipher) as (url, _state):
        assert _pull(url, key, tmp_path / "out.gz", token="wrong") == 1


def test_checksum_mismatch_exits_1(tmp_path):
    identity = pyrage.x25519.Identity.generate()
    key = tmp_path / "key.txt"
    key.write_text(str(identity) + "\n")
    manifest, cipher = _pack(identity)
    manifest["file"]["sha256"] = "0" * 64
    out = tmp_path / "out.gz"

    with _serve(manifest, cipher) as (url, _state):
        assert _pull(url, key, out) == 1
    assert not out.exists()


def test_foreign_identity_exits_1(tmp_path):
    owner = pyrage.x25519.Identity.generate()
    manifest, cipher = _pack(owner)
    key = _identity(tmp_path)  # a different key
    out = tmp_path / "out.gz"

    with _serve(manifest, cipher) as (url, _state):
        assert _pull(url, key, out) == 1
    assert not out.exists()


def test_keygen_writes_a_usable_identity(tmp_path):
    cli = _cli()
    assert cli.main(["keygen", "--out", str(tmp_path)]) == 0
    key = tmp_path / "key.txt"
    assert key.read_text().startswith("AGE-SECRET-KEY-")
    assert key.stat().st_mode & 0o777 == 0o600
    # Refuses to clobber an existing key unless forced.
    assert cli.main(["keygen", "--out", str(tmp_path)]) == 1