"""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." "" " Real QPUs differ." "ignored body" ), "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 = "beforemiddlebodyafter" 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