ai-tax-agent/tests/integration/coverage/test_check_document_coverage_happy_path.py

"""Integration tests for document coverage checking - happy path scenarios."""

# FILE: tests/integration/coverage/test_check_document_coverage_happy_path.py

from datetime import datetime
from unittest.mock import AsyncMock

import pytest

from libs.coverage.evaluator import CoverageEvaluator
from libs.schemas import (
    CompiledCoveragePolicy,
    CoveragePolicy,
    Defaults,
    EvidenceItem,
    OverallStatus,
    Role,
    SchedulePolicy,
    Status,
    StatusClassifier,
    StatusClassifierConfig,
    TaxYearBoundary,
    Trigger,
)


class TestCoverageHappyPath:
    """Test coverage evaluation happy path scenarios"""

    @pytest.fixture
    def mock_kg_client(self):
        """Create mock KG client"""
        client = AsyncMock()

        # Mock successful evidence finding
        client.run_query = AsyncMock(
            return_value=[
                {
                    "doc_id": "DOC-P60-001",
                    "kind": "P60",
                    "page": 1,
                    "bbox": {"x": 100, "y": 200, "width": 300, "height": 50},
                    "ocr_confidence": 0.95,
                    "extract_confidence": 0.92,
                    "date": "2024-05-15",
                }
            ]
        )

        return client

    @pytest.fixture
    def mock_rag_client(self):
        """Create mock RAG client"""
        return AsyncMock()

    @pytest.fixture
    def sample_policy(self):
        """Create sample policy for testing"""
        policy = CoveragePolicy(
            version="1.0",
            jurisdiction="UK",
            tax_year="2024-25",
            tax_year_boundary=TaxYearBoundary(start="2024-04-06", end="2025-04-05"),
            defaults=Defaults(
                confidence_thresholds={"ocr": 0.82, "extract": 0.85},
                date_tolerance_days=30,
            ),
            document_kinds=["P60", "P11D"],
            triggers={
                "SA102": Trigger(any_of=["exists(IncomeItem[type='Employment'])"])
            },
            schedules={
                "SA102": SchedulePolicy(
                    evidence=[
                        EvidenceItem(
                            id="P60",
                            role=Role.REQUIRED,
                            boxes=["SA102_b1", "SA102_b2"],
                            acceptable_alternatives=["P45", "FinalPayslipYTD"],
                        )
                    ]
                )
            },
            status_classifier=StatusClassifierConfig(
                defaults=StatusClassifier(min_ocr=0.82, min_extract=0.85),
                present_verified=StatusClassifier(min_ocr=0.82, min_extract=0.85),
                present_unverified=StatusClassifier(min_ocr=0.60, min_extract=0.70),
                conflicting=StatusClassifier(min_ocr=0.60, min_extract=0.70),
            ),
            conflict_resolution={"precedence": ["P60"]},
            question_templates={"default": {"text": "test", "why": "test"}},
            privacy={},
        )

        # Create compiled policy with mock predicates
        compiled = CompiledCoveragePolicy(
            policy=policy,
            compiled_predicates={
                "exists(IncomeItem[type='Employment'])": lambda tid, ty: True  # Always true for test
            },
            compiled_at=datetime.utcnow(),
            hash="test-hash",
            source_files=["test.yaml"],
        )

        return compiled

    @pytest.mark.asyncio
    async def test_complete_coverage_happy_path(
        self, mock_kg_client, mock_rag_client, sample_policy
    ):
        """Test complete coverage evaluation with all evidence present"""
        evaluator = CoverageEvaluator(
            kg_client=mock_kg_client, rag_client=mock_rag_client
        )

        report = await evaluator.check_document_coverage(
            taxpayer_id="T-001",
            tax_year="2024-25",
            policy=sample_policy,
        )

        # Should have identified SA102 as required
        assert "SA102" in report.schedules_required

        # Should have overall OK status
        assert report.overall_status == OverallStatus.OK

        # Should have no blocking items
        assert len(report.blocking_items) == 0

        # Should have coverage for SA102
        sa102_coverage = next(c for c in report.coverage if c.schedule_id == "SA102")
        assert sa102_coverage.status == OverallStatus.OK

        # Should have P60 evidence marked as verified
        p60_evidence = next(e for e in sa102_coverage.evidence if e.id == "P60")
        assert p60_evidence.status == Status.PRESENT_VERIFIED
        assert len(p60_evidence.found) == 1
        assert p60_evidence.found[0].doc_id == "DOC-P60-001"

    @pytest.mark.asyncio
    async def test_infer_required_schedules(self, mock_kg_client, sample_policy):
        """Test schedule inference based on triggers"""
        evaluator = CoverageEvaluator(kg_client=mock_kg_client)

        required = await evaluator.infer_required_schedules(
            taxpayer_id="T-001",
            tax_year="2024-25",
            policy=sample_policy,
        )

        # Should identify SA102 as required (predicate returns True)
        assert "SA102" in required

    @pytest.mark.asyncio
    async def test_find_evidence_docs(self, mock_kg_client, sample_policy):
        """Test evidence document finding"""
        evaluator = CoverageEvaluator(kg_client=mock_kg_client)

        evidence_map = await evaluator.find_evidence_docs(
            taxpayer_id="T-001",
            tax_year="2024-25",
            evidence_ids=["P60"],
            policy=sample_policy,
        )

        # Should find P60 evidence
        assert "P60" in evidence_map
        assert len(evidence_map["P60"]) == 1

        found_evidence = evidence_map["P60"][0]
        assert found_evidence.doc_id == "DOC-P60-001"
        assert found_evidence.kind == "P60"
        assert found_evidence.ocr_confidence == 0.95

    @pytest.mark.asyncio
    async def test_build_reason_and_citations(
        self, mock_kg_client, mock_rag_client, sample_policy
    ):
        """Test reason and citation building"""
        evaluator = CoverageEvaluator(
            kg_client=mock_kg_client, rag_client=mock_rag_client
        )

        # Mock KG citations
        mock_kg_client.run_query.return_value = [
            {
                "rule_id": "UK.SA102.P60.Required",
                "doc_id": "SA102-Notes-2025",
                "locator": "p.3 §1.1",
            }
        ]

        evidence_item = sample_policy.policy.schedules["SA102"].evidence[0]

        reason, citations = await evaluator.build_reason_and_citations(
            schedule_id="SA102",
            evidence_item=evidence_item,
            status=Status.PRESENT_VERIFIED,
            taxpayer_id="T-001",
            tax_year="2024-25",
            policy=sample_policy,
        )

        # Should build appropriate reason
        assert "P60" in reason
        assert "verified" in reason.lower()

        # Should have citations
        assert len(citations) > 0

    @pytest.mark.asyncio
    async def test_multiple_schedules_coverage(self, mock_kg_client, sample_policy):
        """Test coverage evaluation with multiple schedules"""
        # Add another schedule to policy
        sample_policy.policy.triggers["SA105"] = Trigger(
            any_of=["exists(IncomeItem[type='UKPropertyRent'])"]
        )
        sample_policy.policy.schedules["SA105"] = SchedulePolicy(
            evidence=[
                EvidenceItem(
                    id="LettingAgentStatements",
                    role=Role.REQUIRED,
                    boxes=["SA105_b5"],
                )
            ]
        )

        # Add predicate for SA105
        sample_policy.compiled_predicates[
            "exists(IncomeItem[type='UKPropertyRent'])"
        ] = lambda tid, ty: True

        # Mock evidence for both schedules
        def mock_query_side_effect(query, params):
            if "P60" in params.get("kinds", []):
                return [
                    {
                        "doc_id": "DOC-P60-001",
                        "kind": "P60",
                        "page": 1,
                        "bbox": {},
                        "ocr_confidence": 0.95,
                        "extract_confidence": 0.92,
                        "date": "2024-05-15",
                    }
                ]
            elif "LettingAgentStatements" in params.get("kinds", []):
                return [
                    {
                        "doc_id": "DOC-AGENT-001",
                        "kind": "LettingAgentStatements",
                        "page": 1,
                        "bbox": {},
                        "ocr_confidence": 0.88,
                        "extract_confidence": 0.90,
                        "date": "2024-06-01",
                    }
                ]
            return []

        mock_kg_client.run_query.side_effect = mock_query_side_effect

        evaluator = CoverageEvaluator(kg_client=mock_kg_client)

        report = await evaluator.check_document_coverage(
            taxpayer_id="T-001",
            tax_year="2024-25",
            policy=sample_policy,
        )

        # Should identify both schedules as required
        assert "SA102" in report.schedules_required
        assert "SA105" in report.schedules_required

        # Should have coverage for both schedules
        assert len(report.coverage) == 2

        # Both should be OK
        assert report.overall_status == OverallStatus.OK

    @pytest.mark.asyncio
    async def test_conditional_evidence_not_required(
        self, mock_kg_client, sample_policy
    ):
        """Test that conditional evidence is skipped when condition not met"""
        # Add conditional evidence to SA102
        conditional_evidence = EvidenceItem(
            id="P11D",
            role=Role.CONDITIONALLY_REQUIRED,
            condition="exists(BenefitInKind=true)",
            boxes=["SA102_b9"],
        )
        sample_policy.policy.schedules["SA102"].evidence.append(conditional_evidence)

        # Add predicate that returns False (condition not met)
        sample_policy.compiled_predicates["exists(BenefitInKind=true)"] = (
            lambda tid, ty: False
        )

        evaluator = CoverageEvaluator(kg_client=mock_kg_client)

        report = await evaluator.check_document_coverage(
            taxpayer_id="T-001",
            tax_year="2024-25",
            policy=sample_policy,
        )

        # Should have SA102 coverage
        sa102_coverage = next(c for c in report.coverage if c.schedule_id == "SA102")

        # Should only have P60 evidence (P11D should be skipped)
        evidence_ids = [e.id for e in sa102_coverage.evidence]
        assert "P60" in evidence_ids
        assert "P11D" not in evidence_ids

    @pytest.mark.asyncio
    async def test_conditional_evidence_required(self, mock_kg_client, sample_policy):
        """Test that conditional evidence is included when condition is met"""
        # Add conditional evidence to SA102
        conditional_evidence = EvidenceItem(
            id="P11D",
            role=Role.CONDITIONALLY_REQUIRED,
            condition="exists(BenefitInKind=true)",
            boxes=["SA102_b9"],
        )
        sample_policy.policy.schedules["SA102"].evidence.append(conditional_evidence)

        # Add predicate that returns True (condition met)
        sample_policy.compiled_predicates["exists(BenefitInKind=true)"] = (
            lambda tid, ty: True
        )

        # Mock evidence finding for P11D
        def mock_query_side_effect(query, params):
            if "P60" in params.get("kinds", []):
                return [
                    {
                        "doc_id": "DOC-P60-001",
                        "kind": "P60",
                        "page": 1,
                        "bbox": {},
                        "ocr_confidence": 0.95,
                        "extract_confidence": 0.92,
                        "date": "2024-05-15",
                    }
                ]
            elif "P11D" in params.get("kinds", []):
                return [
                    {
                        "doc_id": "DOC-P11D-001",
                        "kind": "P11D",
                        "page": 1,
                        "bbox": {},
                        "ocr_confidence": 0.90,
                        "extract_confidence": 0.88,
                        "date": "2024-07-06",
                    }
                ]
            return []

        mock_kg_client.run_query.side_effect = mock_query_side_effect

        evaluator = CoverageEvaluator(kg_client=mock_kg_client)

        report = await evaluator.check_document_coverage(
            taxpayer_id="T-001",
            tax_year="2024-25",
            policy=sample_policy,
        )

        # Should have SA102 coverage
        sa102_coverage = next(c for c in report.coverage if c.schedule_id == "SA102")

        # Should have both P60 and P11D evidence
        evidence_ids = [e.id for e in sa102_coverage.evidence]
        assert "P60" in evidence_ids
        assert "P11D" in evidence_ids

        # Both should be verified
        p60_evidence = next(e for e in sa102_coverage.evidence if e.id == "P60")
        p11d_evidence = next(e for e in sa102_coverage.evidence if e.id == "P11D")
        assert p60_evidence.status == Status.PRESENT_VERIFIED
        assert p11d_evidence.status == Status.PRESENT_VERIFIED

    @pytest.mark.asyncio
    async def test_no_schedules_required(self, mock_kg_client, sample_policy):
        """Test coverage when no schedules are required"""
        # Make predicate return False (no employment income)
        sample_policy.compiled_predicates["exists(IncomeItem[type='Employment'])"] = (
            lambda tid, ty: False
        )

        evaluator = CoverageEvaluator(kg_client=mock_kg_client)

        report = await evaluator.check_document_coverage(
            taxpayer_id="T-001",
            tax_year="2024-25",
            policy=sample_policy,
        )

        # Should have no required schedules
        assert len(report.schedules_required) == 0

        # Should have OK status (nothing required, nothing missing)
        assert report.overall_status == OverallStatus.OK

        # Should have no coverage items
        assert len(report.coverage) == 0

        # Should have no blocking items
        assert len(report.blocking_items) == 0