arborist/tests/test_mesh_cli_pull.py

"""CLI integration: `aborist mesh pull`.

Closes the request half of the gossip loop. Wire-layer Merkle / signature
verification is covered in test_mesh_wire_e2e.py — these tests cover the
argparse plumbing, gating, and the local-ingest round-trip that the pull
verb performs after the wire client returns a verified body.
"""

from __future__ import annotations

import json
from pathlib import Path

from aborist.cli import build_parser
from aborist.document import Document
from aborist.ingest import ingest_source
from aborist.mesh import init_identity, set_enabled
from aborist.mesh.members import add_member
from aborist.mesh.wire import MeshWireServer
from aborist.store import connect


class _OneDocSource:
    source_type = "mesh_pull_test"

    def __init__(self, uri: str, content: str, title: str):
        self._doc = Document(
            uri=uri, content=content, source_type=self.source_type, title=title
        )

    def iter_documents(self):
        yield self._doc


def _ingest(db_path, *, uri: str, content: str, title: str) -> str:
    conn = connect(db_path)
    try:
        ingest_source(conn, _OneDocSource(uri, content, title))
        row = conn.execute(
            "SELECT document_root FROM documents WHERE document_uri=?", (uri,)
        ).fetchone()
    finally:
        conn.close()
    return row["document_root"]


def _bootstrap_two_peers(tmp_path: Path):
    alice_db = tmp_path / "alice.db"
    bob_db = tmp_path / "bob.db"
    a_conn = connect(alice_db)
    try:
        alice = init_identity(a_conn, group_name="t", member_id="alice")
        set_enabled(a_conn, True)
    finally:
        a_conn.close()
    b_conn = connect(bob_db)
    try:
        bob = init_identity(b_conn, group_name="t", member_id="bob")
        set_enabled(b_conn, True)
    finally:
        b_conn.close()
    a_conn = connect(alice_db)
    try:
        add_member(a_conn, member_id="bob", sign_pub=bob.sign_pub, dh_pub=bob.dh_pub)
    finally:
        a_conn.close()
    b_conn = connect(bob_db)
    try:
        add_member(b_conn, member_id="alice", sign_pub=alice.sign_pub, dh_pub=alice.dh_pub)
    finally:
        b_conn.close()
    return alice_db, bob_db, alice, bob


# ---------------------------------------------------------------------------
# Gating: mesh must be initialized AND enabled
# ---------------------------------------------------------------------------


def test_pull_refuses_when_mesh_not_initialized(tmp_path, capsys):
    db = tmp_path / "x.db"
    parser = build_parser()
    args = parser.parse_args([
        "--db", str(db),
        "mesh", "pull",
        "--root", "ab" * 32,
        "--peer", "http://0:0",
    ])
    rc = args.func(args)
    assert rc == 2
    assert "not initialized" in capsys.readouterr().err


def test_pull_refuses_when_mesh_disabled(tmp_path, capsys):
    db = tmp_path / "x.db"
    conn = connect(db)
    try:
        init_identity(conn, group_name="t", member_id="alice")
        # No set_enabled — flag stays off.
    finally:
        conn.close()
    parser = build_parser()
    args = parser.parse_args([
        "--db", str(db),
        "mesh", "pull",
        "--root", "ab" * 32,
        "--peer", "http://0:0",
    ])
    rc = args.func(args)
    assert rc == 2
    assert "enabled" in capsys.readouterr().err


# ---------------------------------------------------------------------------
# Happy path: bob pulls a doc alice has
# ---------------------------------------------------------------------------


def test_pull_fetches_and_ingests_doc_from_peer(tmp_path, capsys):
    """Alice has a document; bob pulls it. Bob's local DB ends up with the
    same document_root, an audit event is written, and the response JSON
    reports `pulled`."""
    alice_db, bob_db, _alice, _bob = _bootstrap_two_peers(tmp_path)

    # Alice ingests; bob does not have it yet.
    document_root = _ingest(
        alice_db,
        uri="https://example.com/alpha",
        content="Alpha content. The quick brown fox jumps over the lazy dog.",
        title="Alpha",
    )

    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    srv.start_in_thread()
    try:
        parser = build_parser()
        args = parser.parse_args([
            "--db", str(bob_db),
            "mesh", "pull",
            "--root", document_root,
            "--peer", srv.url,
        ])
        rc = args.func(args)
        assert rc == 0
    finally:
        srv.stop()

    payload = json.loads(capsys.readouterr().out)
    assert payload["status"] == "pulled"
    assert payload["document_root"] == document_root
    assert payload["document_uri"] == "https://example.com/alpha"
    assert payload["shard"] == str(bob_db)
    assert isinstance(payload["audit_event_hash"], str)
    assert len(payload["audit_event_hash"]) == 64

    # Bob's DB now contains the document at the same root.
    b_conn = connect(bob_db)
    try:
        row = b_conn.execute(
            "SELECT document_root, document_uri FROM documents WHERE document_root=?",
            (document_root,),
        ).fetchone()
        audit = b_conn.execute(
            "SELECT event_type, subject_root, body FROM audit_events "
            "WHERE event_type='mesh_pulled' ORDER BY seq DESC LIMIT 1"
        ).fetchone()
    finally:
        b_conn.close()

    assert row is not None
    assert row["document_uri"] == "https://example.com/alpha"
    assert audit is not None
    assert audit["subject_root"] == document_root
    body = json.loads(audit["body"])
    assert body["document_root"] == document_root
    assert body["document_uri"] == "https://example.com/alpha"
    assert body["peer"] == srv.url


# ---------------------------------------------------------------------------
# Idempotence: pulling a doc bob already has returns `already_present`
# ---------------------------------------------------------------------------


def test_pull_already_present_when_local_has_doc(tmp_path, capsys):
    alice_db, bob_db, _alice, _bob = _bootstrap_two_peers(tmp_path)

    # Both peers ingest the same content -> identical document_root.
    content = "Beta content. The quick brown fox jumps over the lazy dog."
    a_root = _ingest(alice_db, uri="https://example.com/beta", content=content, title="Beta")
    b_root = _ingest(bob_db, uri="https://example.com/beta", content=content, title="Beta")
    assert a_root == b_root

    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    srv.start_in_thread()
    try:
        parser = build_parser()
        args = parser.parse_args([
            "--db", str(bob_db),
            "mesh", "pull",
            "--root", a_root,
            "--peer", srv.url,
        ])
        rc = args.func(args)
        assert rc == 0
    finally:
        srv.stop()

    payload = json.loads(capsys.readouterr().out)
    assert payload["status"] == "already_present"
    assert payload["document_root"] == a_root
    assert payload["document_uri"] == "https://example.com/beta"

    # No mesh_pulled event was written since we short-circuited.
    b_conn = connect(bob_db)
    try:
        rows = b_conn.execute(
            "SELECT COUNT(*) FROM audit_events WHERE event_type='mesh_pulled'"
        ).fetchone()[0]
    finally:
        b_conn.close()
    assert rows == 0


# ---------------------------------------------------------------------------
# Failure: peer returns 404 for unknown root
# ---------------------------------------------------------------------------


def test_pull_returns_rc2_on_unknown_root(tmp_path, capsys):
    """Peer responds 404 (no such document_root) -> rc=2, error to stderr."""
    alice_db, bob_db, _alice, _bob = _bootstrap_two_peers(tmp_path)

    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    srv.start_in_thread()
    try:
        parser = build_parser()
        args = parser.parse_args([
            "--db", str(bob_db),
            "mesh", "pull",
            "--root", "ff" * 32,
            "--peer", srv.url,
        ])
        rc = args.func(args)
        assert rc == 2
    finally:
        srv.stop()

    err = capsys.readouterr().err
    assert "pull failed" in err