tests/test_knowledge_extractor.py

"""
GovBridge India — Unit Tests for Knowledge Extraction Pipeline
Sprint 30: PROJECT INDRA Phase 1.5

Tests cover:
  1. Prompt construction
  2. Pydantic validation (valid predicates, invalid predicates)
  3. Text chunking with overlap
  4. Triplet deduplication
  5. Entity normalization

Run: pytest tests/test_knowledge_extractor.py -v
"""

import sys
import os
import pytest

# Ensure gov_backend and ingestion are on sys.path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ingestion"))


# ═══════════════════════════════════════════════════════════════
# TEST 1: Predicate Ontology
# ═══════════════════════════════════════════════════════════════

class TestGovPredicate:
    def test_all_10_predicates_exist(self):
        from prompts import VALID_PREDICATES
        expected = {
            "AMENDS", "SUPERSEDES", "REPEALS", "ENACTS",
            "DELEGATES_TO", "FUNDS", "MANDATES", "APPLIES_TO",
            "DEFINES", "PENALIZES"
        }
        assert VALID_PREDICATES == expected

    def test_predicate_enum_values(self):
        from prompts import GovPredicate
        assert GovPredicate.AMENDS.value == "AMENDS"
        assert GovPredicate.DELEGATES_TO.value == "DELEGATES_TO"


# ═══════════════════════════════════════════════════════════════
# TEST 2: Pydantic Validation
# ═══════════════════════════════════════════════════════════════

class TestKnowledgeTriplet:
    def test_valid_triplet(self):
        from knowledge_extractor import KnowledgeTriplet
        t = KnowledgeTriplet(
            subject="Finance Act 2024",
            predicate="AMENDS",
            object="Income Tax Act 1961"
        )
        assert t.predicate == "AMENDS"
        assert t.subject == "Finance Act 2024"

    def test_predicate_normalized_to_uppercase(self):
        from knowledge_extractor import KnowledgeTriplet
        t = KnowledgeTriplet(
            subject="Act A",
            predicate="amends",  # lowercase
            object="Act B"
        )
        assert t.predicate == "AMENDS"

    def test_invalid_predicate_rejected(self):
        from knowledge_extractor import KnowledgeTriplet
        from pydantic import ValidationError
        with pytest.raises(ValidationError):
            KnowledgeTriplet(
                subject="Act A",
                predicate="MODIFIES",  # Not in ontology
                object="Act B"
            )

    def test_empty_subject_rejected(self):
        from knowledge_extractor import KnowledgeTriplet
        from pydantic import ValidationError
        with pytest.raises(ValidationError):
            KnowledgeTriplet(
                subject="",
                predicate="AMENDS",
                object="Act B"
            )

    def test_whitespace_trimmed(self):
        from knowledge_extractor import KnowledgeTriplet
        t = KnowledgeTriplet(
            subject="  Finance Act 2024  ",
            predicate="  REPEALS  ",
            object="  Old Act  "
        )
        assert t.subject == "Finance Act 2024"
        assert t.object == "Old Act"


# ═══════════════════════════════════════════════════════════════
# TEST 3: Text Chunking
# ═══════════════════════════════════════════════════════════════

class TestTextChunking:
    def test_short_text_single_chunk(self):
        from knowledge_extractor import chunk_text_for_extraction
        chunks = chunk_text_for_extraction("Short text", chunk_size=1000)
        assert len(chunks) == 1
        assert chunks[0] == "Short text"

    def test_long_text_produces_multiple_chunks(self):
        from knowledge_extractor import chunk_text_for_extraction
        # Create text that's clearly > 1000 chars
        text = "A" * 500 + "\n\n" + "B" * 500 + "\n\n" + "C" * 500
        chunks = chunk_text_for_extraction(text, chunk_size=600, overlap=100)
        assert len(chunks) >= 2

    def test_overlap_is_applied(self):
        from knowledge_extractor import chunk_text_for_extraction
        # Create text with clear paragraph boundaries
        text = ("Para one. " * 50) + "\n\n" + ("Para two. " * 50) + "\n\n" + ("Para three. " * 50)
        chunks = chunk_text_for_extraction(text, chunk_size=300, overlap=50)
        # Verify overlap: last N chars of chunk[0] should appear in chunk[1]
        if len(chunks) > 1:
            tail = chunks[0][-50:]
            assert tail in chunks[1] or len(chunks[1]) > 50


# ═══════════════════════════════════════════════════════════════
# TEST 4: Deduplication
# ═══════════════════════════════════════════════════════════════

class TestDeduplication:
    def test_exact_duplicates_removed(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
            KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 1

    def test_case_insensitive_dedup(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="Finance Act", predicate="AMENDS", object="IT Act"),
            KnowledgeTriplet(subject="finance act", predicate="AMENDS", object="it act"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 1

    def test_different_predicates_kept(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="Act A", predicate="AMENDS", object="Act B"),
            KnowledgeTriplet(subject="Act A", predicate="SUPERSEDES", object="Act B"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 2

    def test_article_noise_dedup(self):
        from knowledge_extractor import KnowledgeTriplet, deduplicate_triplets
        triplets = [
            KnowledgeTriplet(subject="The Finance Act", predicate="AMENDS", object="The IT Act"),
            KnowledgeTriplet(subject="Finance Act", predicate="AMENDS", object="IT Act"),
        ]
        result = deduplicate_triplets(triplets)
        assert len(result) == 1


# ═══════════════════════════════════════════════════════════════
# TEST 5: Entity Normalization
# ═══════════════════════════════════════════════════════════════

class TestEntityNormalization:
    def test_normalize_strips_whitespace(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("  Finance Act  ") == "finance act"

    def test_normalize_removes_articles(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("The Finance Act") == "finance act"

    def test_normalize_collapses_spaces(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("Finance   Act   2024") == "finance act 2024"

    def test_normalize_removes_legal_filler(self):
        from knowledge_extractor import _normalize_entity
        assert _normalize_entity("said Ministry of Finance") == "ministry of finance"


# ═══════════════════════════════════════════════════════════════
# TEST 6: Prompt Construction
# ═══════════════════════════════════════════════════════════════

class TestPromptConstruction:
    def test_build_chunk_user_prompt_includes_context(self):
        from prompts import build_chunk_user_prompt
        ctx = {
            "document_title": "Finance Act 2024",
            "primary_body": "Ministry of Finance",
            "key_definitions": ["tax", "income"],
            "document_date": "2024-01-15"
        }
        prompt = build_chunk_user_prompt(ctx, "Some chunk text here")
        assert "Finance Act 2024" in prompt
        assert "Ministry of Finance" in prompt
        assert "Some chunk text here" in prompt
        assert "tax, income" in prompt

    def test_build_chunk_user_prompt_handles_empty_context(self):
        from prompts import build_chunk_user_prompt
        ctx = {}
        prompt = build_chunk_user_prompt(ctx, "Chunk text")
        assert "Unknown" in prompt
        assert "Chunk text" in prompt