arborist/tests/test_concepts_extract.py

"""Tests for ``arborist.concepts.extract`` — concept-relation
extractors that derive synonyms / IDF / FTS5-titles from corpus
state. Module had zero direct tests despite being the write-side
of every concept-relations row the existing test_concepts.py
exercises on the read-side.

Coverage:
  - _title_tokens (pure private helper) — stopword strip,
    underscore→space, length≥4 filter, lowercase, hyphen-handling
  - EXTRACTORS registry contract — keys + signature
  - backfill_documents_fts on a synthetic mini-shard
  - link_reciprocity_synonym idempotency on a 2-doc reciprocal pair

The extractors operate on real shard schema; we build a minimal
in-memory schema that matches store.py's contract just enough to
exercise the code paths.
"""

from __future__ import annotations

import sqlite3

import pytest

from arborist.concepts.extract import (
    EXTRACTORS,
    _title_tokens,
    backfill_documents_fts,
    backfill_token_idf,
    link_reciprocity_synonym,
)
from arborist.store import SCHEMA_SQL


# --- _title_tokens (pure) ------------------------------------------


def test_title_tokens_empty_string():
    assert _title_tokens("") == set()


def test_title_tokens_strips_stopwords():
    """Documented stopwords (the/a/an/of/and/...) drop out."""
    out = _title_tokens("the history of computing")
    assert "the" not in out  # stopword
    assert "history" in out
    assert "computing" in out
    # "of" stopword + length<4 anyway
    assert "of" not in out


def test_title_tokens_underscore_to_space():
    """Wikipedia title style 'Foo_Bar' → split on underscore."""
    out = _title_tokens("New_York_City")
    assert "york" in out
    assert "city" in out


def test_title_tokens_length_four_floor():
    """Tokens length<4 dropped (3-letter words = noisy anchors)."""
    out = _title_tokens("the cat ran")
    assert "cat" not in out  # 3 chars
    assert "ran" not in out  # 3 chars


def test_title_tokens_lowercase():
    """Output always lowercase regardless of input case."""
    out = _title_tokens("Quantum CHROMODYNAMICS")
    assert "quantum" in out
    assert "chromodynamics" in out


def test_title_tokens_handles_hyphenated_token():
    """Title-token regex `[a-z][a-z0-9'\\-]+` keeps hyphens inside
    a token. 'state-of-the-art' → 'state-of-the-art' as one token
    (length>=4)."""
    out = _title_tokens("state-of-the-art technology")
    # The regex starts on a letter then captures alphanum/apos/hyphen.
    # 'state-of-the-art' becomes one token.
    assert any("-" in t for t in out)
    assert "technology" in out


def test_title_tokens_dedupe():
    """Output is a set — duplicates collapse."""
    out = _title_tokens("Quantum quantum mechanics")
    quantum_count = sum(1 for t in out if t == "quantum")
    assert quantum_count == 1


def test_title_tokens_skips_pure_digits():
    """Token must START with [a-z] per the regex; 1984 → no match."""
    out = _title_tokens("1984 novel")
    assert "1984" not in out
    assert "novel" in out


# --- EXTRACTORS registry --------------------------------------------


def test_extractors_registry_has_documented_keys():
    """Three extractors documented in the module-level docstring."""
    assert "link_reciprocity" in EXTRACTORS
    assert "token_idf" in EXTRACTORS
    assert "documents_fts" in EXTRACTORS


def test_extractors_registry_callables():
    """Each registry value is callable."""
    for name, fn in EXTRACTORS.items():
        assert callable(fn), f"EXTRACTORS[{name!r}] is not callable"


def test_extractors_registry_link_reciprocity_routes_to_function():
    """The dispatch table maps to the actual implementations."""
    assert EXTRACTORS["link_reciprocity"] is link_reciprocity_synonym
    assert EXTRACTORS["token_idf"] is backfill_token_idf
    assert EXTRACTORS["documents_fts"] is backfill_documents_fts


# --- backfill_documents_fts ----------------------------------------


@pytest.fixture
def empty_shard(tmp_path):
    """Fresh SQLite with arborist's full schema applied."""
    db_path = tmp_path / "shard.db"
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA_SQL)
    yield conn
    conn.close()


def _insert_doc(conn, root: str, uri: str, title: str):
    conn.execute(
        "INSERT INTO documents "
        "(document_root, document_uri, source_type, kind, "
        " compression_depth, title, chunking_version, "
        " canonicalization_version, schema_version, ingest_ts) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        (root, uri, "html", "surface", 0, title,
         "tok-512-v1", "norm-v1", "v9.8.0", 0),
    )


def test_backfill_documents_fts_indexes_titles(empty_shard):
    _insert_doc(empty_shard, "a" * 64, "uri-a", "Apple Inc")
    _insert_doc(empty_shard, "b" * 64, "uri-b", "Banana fruit")
    _insert_doc(empty_shard, "c" * 64, "uri-c", "Cherry tree")
    empty_shard.commit()

    result = backfill_documents_fts(empty_shard)
    assert result["rows_indexed"] == 3
    assert result["elapsed_ms"] >= 0


def test_backfill_documents_fts_idempotent(empty_shard):
    """Re-running on the same shard yields identical row count
    (DELETE + INSERT inside)."""
    _insert_doc(empty_shard, "a" * 64, "uri-a", "First")
    _insert_doc(empty_shard, "b" * 64, "uri-b", "Second")
    empty_shard.commit()

    r1 = backfill_documents_fts(empty_shard)
    r2 = backfill_documents_fts(empty_shard)
    assert r1["rows_indexed"] == r2["rows_indexed"] == 2


def test_backfill_documents_fts_skips_null_titles(empty_shard):
    """Documents with NULL title are not indexed (no FTS5 row)."""
    _insert_doc(empty_shard, "a" * 64, "uri-a", "Has Title")
    # Insert a row with no title.
    empty_shard.execute(
        "INSERT INTO documents "
        "(document_root, document_uri, source_type, kind, "
        " compression_depth, title, chunking_version, "
        " canonicalization_version, schema_version, ingest_ts) "
        "VALUES (?, ?, ?, ?, ?, NULL, ?, ?, ?, ?)",
        ("b" * 64, "uri-b", "html", "surface", 0,
         "tok-512-v1", "norm-v1", "v9.8.0", 0),
    )
    empty_shard.commit()
    result = backfill_documents_fts(empty_shard)
    assert result["rows_indexed"] == 1  # only the titled doc


def test_backfill_documents_fts_default_derived_from_doesnt_break(empty_shard):
    """Default derived_from kwarg path works (None → derived from
    the function name internally)."""
    _insert_doc(empty_shard, "a" * 64, "uri-a", "Test")
    empty_shard.commit()
    # Default kwarg path
    r1 = backfill_documents_fts(empty_shard)
    # Explicit kwarg path
    r2 = backfill_documents_fts(empty_shard, derived_from="custom-extractor")
    assert r1["rows_indexed"] == r2["rows_indexed"]


# --- link_reciprocity_synonym --------------------------------------


def _insert_edge(conn, src: str, dst: str, dst_root: str):
    conn.execute(
        "INSERT INTO edges "
        "(src_root, edge_type, dst_root, dst_uri, anchor) "
        "VALUES (?, 'related_doc', ?, ?, '')",
        (src, dst_root, dst),
    )


def test_link_reciprocity_no_edges_returns_zero(empty_shard):
    result = link_reciprocity_synonym(empty_shard)
    assert result["reciprocal_pairs"] == 0
    assert result["synonyms_inserted"] == 0


def test_link_reciprocity_one_way_edge_no_pair(empty_shard):
    """Edge A→B without reverse B→A → no reciprocal pair → no
    synonym."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple Inc")
    _insert_doc(empty_shard, b, "uri-b", "Banana fruit")
    _insert_edge(empty_shard, a, "uri-b", b)  # only A→B
    empty_shard.commit()
    result = link_reciprocity_synonym(empty_shard)
    assert result["reciprocal_pairs"] == 0


def test_link_reciprocity_pair_creates_synonyms(empty_shard):
    """A↔B with title tokens → synonym edges."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple computing")
    _insert_doc(empty_shard, b, "uri-b", "Banana technology")
    _insert_edge(empty_shard, a, "uri-b", b)
    _insert_edge(empty_shard, b, "uri-a", a)
    empty_shard.commit()
    result = link_reciprocity_synonym(empty_shard)
    assert result["reciprocal_pairs"] == 1
    # Cross-product of 2 tokens × 2 tokens, minus self-overlap.
    # apple/computing × banana/technology — no overlap → 4 pairs.
    assert result["synonyms_inserted"] == 4


def test_link_reciprocity_idempotent(empty_shard):
    """Re-running adds no new rows (concept_relations PK enforces
    uniqueness)."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple computing")
    _insert_doc(empty_shard, b, "uri-b", "Banana technology")
    _insert_edge(empty_shard, a, "uri-b", b)
    _insert_edge(empty_shard, b, "uri-a", a)
    empty_shard.commit()
    r1 = link_reciprocity_synonym(empty_shard)
    r2 = link_reciprocity_synonym(empty_shard)
    assert r1["synonyms_inserted"] >= 1
    assert r2["synonyms_skipped"] == r1["synonyms_inserted"]
    assert r2["synonyms_inserted"] == 0


def test_link_reciprocity_skips_self_overlap_token(empty_shard):
    """If both titles share a token (e.g. 'computing' in both),
    that's not a synonym (same word) — skipped from the cross
    product."""
    a, b = "a" * 64, "b" * 64
    _insert_doc(empty_shard, a, "uri-a", "Apple computing")
    _insert_doc(empty_shard, b, "uri-b", "Banana computing")
    _insert_edge(empty_shard, a, "uri-b", b)
    _insert_edge(empty_shard, b, "uri-a", a)
    empty_shard.commit()
    result = link_reciprocity_synonym(empty_shard)
    # 4 cross pairs minus the (computing, computing) self-overlap = 3
    # apple→banana, apple→computing, banana→computing, computing→apple,
    # computing→banana. Self-pair (computing, computing) excluded.
    # Implementation: cross = {(a,b) for a in A for b in B if a != b}
    # A = {apple, computing}, B = {banana, computing}
    # → (apple, banana), (apple, computing), (computing, banana) = 3
    assert result["synonyms_inserted"] == 3