arborist/tests/test_grok_source.py

"""Tests for the Grok account-export source.

Synthetic fixture mirrors the real prod-grok-backend.json shape: a
`conversations` array of `{conversation, responses}` entries, plus a
`media_posts` array of generation prompts. Real personal data is NOT
imported into the test suite.
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from arborist.ingest import ingest_source
from arborist.sources.grok import (
    GrokExportSource,
    GrokMediaPostsSource,
    _extract_urls,
    _normalize_ts,
    _scrub,
)
from arborist.store import connect


def _make_export(tmp_path: Path) -> Path:
    """Build a minimal but representative export tree, return its root."""
    user_dir = tmp_path / "ttl" / "30d" / "export_data" / "user-uuid"
    user_dir.mkdir(parents=True)
    payload = {
        "conversations": [
            {
                "conversation": {
                    "id": "conv-aaaa-1111",
                    "title": "Quantum Slop",
                    "create_time": "2026-04-25T12:00:00Z",
                    "user_id": "user-uuid",
                },
                "responses": [
                    {
                        "response": {
                            "_id": "msg-001",
                            "conversation_id": "conv-aaaa-1111",
                            "sender": "human",
                            "model": "grok-4-auto",
                            "create_time": {"$date": "2026-04-25T12:00:00Z"},
                            "message": (
                                "what is up with quantum slop? see "
                                "https://example.com/qs and "
                                "https://example.org/qpu, also (https://example.net/x)."
                            ),
                            "metadata": {},
                        },
                        "share_link": None,
                    },
                    {
                        "response": {
                            "_id": "msg-002",
                            "conversation_id": "conv-aaaa-1111",
                            "sender": "ASSISTANT",
                            "model": "grok-3",
                            # Mongo extended JSON unix millis (real export shape).
                            "create_time": {"$numberLong": "1777306441511"},
                            "message": (
                                "Quantum slop is a misnomer."
                                "<grok:render card_id=\"abc\"/>"
                                " Real QPUs differ.<grok:render card_id=\"d\">"
                                "ignored body</grok:render>"
                            ),
                            "metadata": {},
                        },
                        "share_link": None,
                    },
                ],
            },
            {
                "conversation": {
                    "id": "conv-bbbb-2222",
                    "title": "Empty thread",
                    "create_time": "2026-04-26T00:00:00Z",
                    "user_id": "user-uuid",
                },
                # Both messages empty after scrubbing -> doc should be skipped.
                "responses": [
                    {
                        "response": {
                            "_id": "msg-101",
                            "sender": "human",
                            "message": "",
                            "model": "",
                        },
                        "share_link": None,
                    },
                ],
            },
        ],
        "projects": [],
        "tasks": [],
        "media_posts": [
            {
                "id": "media-7777",
                "user_id": "user-uuid",
                "original_prompt": "a fox standing on a permacomputer",
                "media_type": "image",
                "create_time": "2026-03-12T08:00:00Z",
                "link": "https://assets.example/img/abc.webp",
            },
            {
                "id": "media-8888",
                "user_id": "user-uuid",
                "original_prompt": "",
                "media_type": "image",
                "create_time": "2026-03-12T08:00:01Z",
                "link": "https://assets.example/img/empty.webp",
            },
        ],
    }
    (user_dir / "prod-grok-backend.json").write_text(json.dumps(payload))
    return tmp_path


def test_normalize_ts_handles_iso_string():
    assert _normalize_ts("2026-04-25T12:00:00Z") == "2026-04-25T12:00:00Z"


def test_normalize_ts_handles_mongo_date_wrapper():
    assert _normalize_ts({"$date": "2026-04-25T12:00:00Z"}) == "2026-04-25T12:00:00Z"


def test_normalize_ts_handles_mongo_numberlong_millis():
    # 1_777_306_441_511 ms = 2026-04-27T16:14:01.511Z
    iso = _normalize_ts({"$numberLong": "1777306441511"})
    assert iso.startswith("2026-04-27T16:14:01")
    assert iso.endswith("Z")


def test_normalize_ts_handles_nested_date_numberlong():
    iso = _normalize_ts({"$date": {"$numberLong": "1777306441511"}})
    assert iso.startswith("2026-04-27T16:14:01")


def test_normalize_ts_garbage_returns_empty():
    assert _normalize_ts({"weird": "shape"}) == ""
    assert _normalize_ts(None) == ""


def test_scrub_drops_grok_render_markup():
    s = "before<grok:render card_id=\"x\"/>middle<grok:render card_id=\"y\">body</grok:render>after"
    assert _scrub(s) == "beforemiddleafter"


def test_extract_urls_trims_trailing_punctuation_and_dedupes():
    text = (
        "see https://example.com/qs, and https://example.org/qpu. "
        "also (https://example.net/x). https://example.com/qs again."
    )
    urls = _extract_urls(text)
    assert urls == [
        "https://example.com/qs",
        "https://example.org/qpu",
        "https://example.net/x",
    ]


def test_grok_export_source_yields_one_doc_per_nonempty_conversation(tmp_path):
    root = _make_export(tmp_path)
    src = GrokExportSource(path=root)
    docs = list(src.iter_documents())
    assert len(docs) == 1
    doc = docs[0]
    assert doc.uri == "grok://conversation/conv-aaaa-1111"
    assert doc.source_type == "grok_export"
    assert doc.title == "Quantum Slop"
    # Both turns serialized with sender/model headers.
    assert "human (grok-4-auto)" in doc.content
    assert "assistant (grok-3)" in doc.content
    # Mongo extended JSON shapes must be normalized — no raw wrappers leak.
    assert "$numberLong" not in doc.content
    assert "$date" not in doc.content
    # Render markers stripped.
    assert "grok:render" not in doc.content
    # URLs surfaced as hyperlink edges.
    edge_uris = {e.dst_uri for e in doc.edges}
    assert "https://example.com/qs" in edge_uris
    assert "https://example.org/qpu" in edge_uris
    assert "https://example.net/x" in edge_uris
    assert all(e.edge_type == "hyperlink" for e in doc.edges)


def test_grok_media_source_yields_one_doc_per_nonempty_post(tmp_path):
    root = _make_export(tmp_path)
    src = GrokMediaPostsSource(path=root)
    docs = list(src.iter_documents())
    assert len(docs) == 1
    doc = docs[0]
    assert doc.uri == "grok://media/media-7777"
    assert doc.source_type == "grok_media"
    assert "fox standing on a permacomputer" in doc.content
    asset_edges = [e for e in doc.edges if e.edge_type == "asset"]
    assert len(asset_edges) == 1
    assert asset_edges[0].dst_uri == "https://assets.example/img/abc.webp"


def test_resolves_export_root_from_json_or_dir(tmp_path):
    root = _make_export(tmp_path)
    json_path = (
        root / "ttl" / "30d" / "export_data" / "user-uuid" / "prod-grok-backend.json"
    )
    # Directly the JSON file.
    a = list(GrokExportSource(path=json_path).iter_documents())
    # The user-id directory.
    b = list(GrokExportSource(path=json_path.parent).iter_documents())
    # The export root (auto-walks down).
    c = list(GrokExportSource(path=root).iter_documents())
    assert [d.uri for d in a] == [d.uri for d in b] == [d.uri for d in c]


def test_resolve_missing_path_raises(tmp_path):
    with pytest.raises(FileNotFoundError):
        GrokExportSource(path=tmp_path / "nope")


def test_grok_ingest_round_trip(tmp_path):
    """End-to-end: ingest a tiny export into a fresh DB, verify counts."""
    root = _make_export(tmp_path)
    db_path = tmp_path / "arborist.db"
    conn = connect(db_path)
    try:
        conv_stats = ingest_source(conn, GrokExportSource(path=root))
        media_stats = ingest_source(conn, GrokMediaPostsSource(path=root))
    finally:
        conn.close()
    assert conv_stats.inserted == 1
    assert media_stats.inserted == 1
    # Re-ingest is a no-op (idempotent at content-root level).
    conn = connect(db_path)
    try:
        again = ingest_source(conn, GrokExportSource(path=root))
    finally:
        conn.close()
    assert again.inserted == 0
    assert again.skipped_duplicate == 1


def test_serialized_content_is_deterministic(tmp_path):
    """Same export -> same document_root every time. Required for idempotent
    re-ingest and stable cache keys."""
    root = _make_export(tmp_path)
    a = list(GrokExportSource(path=root).iter_documents())
    b = list(GrokExportSource(path=root).iter_documents())
    assert a[0].content == b[0].content