tests/test_extractive.py

"""Tests for the deterministic extractive engine."""

from researchlink.services import extractive as ex

SAMPLE = (
    "Transformers rely on self-attention. The self-attention mechanism scales "
    "quadratically with sequence length. We propose an efficient attention "
    "variant. Our method reduces memory usage significantly. Experiments show "
    "competitive accuracy on language modeling benchmarks."
)


def test_tokenize_drops_stopwords_and_short():
    toks = ex.tokenize("The self-attention is all you need")
    assert "the" not in toks and "is" not in toks
    assert "self-attention" in toks


def test_split_sentences_counts():
    sents = ex.split_sentences(SAMPLE)
    assert len(sents) == 5
    assert all(len(s.split()) >= 4 for s in sents)


def test_split_sentences_empty():
    assert ex.split_sentences("") == []
    assert ex.split_sentences("too short") == []


def test_rank_sentences_returns_subset_in_order():
    ranked = ex.rank_sentences(SAMPLE, k=2)
    assert len(ranked) == 2
    # results are verbatim spans of the source (no fabrication)
    for s in ranked:
        assert s in ex.split_sentences(SAMPLE)
    # order preserved relative to source
    order = [ex.split_sentences(SAMPLE).index(s) for s in ranked]
    assert order == sorted(order)


def test_rank_sentences_deterministic():
    assert ex.rank_sentences(SAMPLE, k=3) == ex.rank_sentences(SAMPLE, k=3)


def test_rank_sentences_query_boost():
    ranked = ex.rank_sentences(SAMPLE, k=1, query="memory usage")
    assert "memory usage" in ranked[0].lower()


def test_rank_sentences_short_text_returns_all():
    text = "One sentence here now. Second sentence follows here."
    assert len(ex.rank_sentences(text, k=5)) == 2


def test_extract_keywords():
    kws = ex.extract_keywords(SAMPLE, k=5)
    assert "attention" in kws or "self-attention" in kws
    assert len(kws) <= 5
    assert kws == ex.extract_keywords(SAMPLE, k=5)  # deterministic


def test_sentences_matching():
    hits = ex.sentences_matching(SAMPLE, ["memory", "quadratically"], k=5)
    assert any("memory" in h.lower() for h in hits)
    assert any("quadratically" in h.lower() for h in hits)