tests/unit/test_agents.py

"""
Unit Tests: Agents
Tests for all 4 agents and the AgentOrchestrator.
Uses mocked Gemini API to test without real API calls.
"""

import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

MOCK_CLASSIFICATION_JSON = json.dumps(
    {
        "category": "safety_protocol",
        "subcategory": "ventilation",
        "confidence": 0.93,
        "reasoning": "Document contains MSHA ventilation requirements and PPE protocols.",
    }
)

MOCK_SAFETY_JSON = json.dumps(
    {
        "score": 82.0,
        "status": "compliant",
        "confidence": 0.88,
        "hazards": [
            {
                "type": "atmospheric",
                "severity": "medium",
                "description": "Potential methane buildup in Section C",
                "regulation": "30 CFR 75.323",
            }
        ],
        "recommendations": ["Install additional methane detectors in Section C"],
        "compliance_details": {
            "msha_compliant": True,
            "osha_compliant": True,
            "dgms_compliant": True,
            "missing_elements": [],
        },
        "summary": "Generally compliant. One atmospheric hazard identified.",
    }
)

MOCK_ENTITIES_JSON = json.dumps(
    {
        "equipment": ["Caterpillar D11", "Joy 12CM15 Continuous Miner"],
        "chemicals": ["methane (CH4)", "carbon monoxide (CO)"],
        "locations": ["Section C", "Main Heading 3"],
        "personnel": ["Safety Officer", "Mine Foreman"],
        "dates": ["2024-01-15", "quarterly"],
        "regulations": ["30 CFR 75.321", "30 CFR 75.323", "MSHA 1910.134"],
    }
)

MOCK_SUMMARY_JSON = json.dumps(
    {
        "summary": "This document outlines safety procedures for underground coal mining operations, covering ventilation standards, PPE requirements, and emergency evacuation protocols.",
        "key_points": [
            "Methane levels must stay below 1% in all working areas",
            "PPE inspection required before each shift",
            "Emergency drills must occur quarterly",
        ],
        "action_items": [
            "Schedule quarterly emergency drill",
            "Replace methane detectors in Section C",
        ],
        "document_purpose": "Define safety standards for underground coal mine personnel.",
        "confidence": 0.91,
    }
)


def make_mock_model(json_response: str):
    """Create a mock Gemini model that returns json_response."""
    mock_response = MagicMock()
    mock_response.text = json_response

    mock_model = MagicMock()
    mock_model.generate_content = MagicMock(return_value=mock_response)
    return mock_model


def make_mock_client(json_response: str):
    """Create a mock OpenAI client that returns json_response."""
    mock_choice = MagicMock()
    mock_choice.message.content = json_response

    mock_completion = MagicMock()
    mock_completion.choices = [mock_choice]

    mock_chat = MagicMock()
    mock_chat.completions.create = AsyncMock(return_value=mock_completion)

    mock_client = MagicMock()
    mock_client.chat = mock_chat
    return mock_client


class TestClassifierAgent:
    """Tests for ClassifierAgent."""

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_returns_expected_keys(self, sample_mining_text):
        """ClassifierAgent.analyze() returns dict with required keys."""
        from app.agents.classifier import ClassifierAgent

        agent = ClassifierAgent()
        agent.client = make_mock_client(MOCK_CLASSIFICATION_JSON)

        result = await agent.analyze(sample_mining_text)

        assert "category" in result
        assert "subcategory" in result
        assert "confidence" in result
        assert "reasoning" in result

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_maps_category_correctly(self, sample_mining_text):
        """ClassifierAgent maps 'safety_protocol' string to correct value."""
        from app.agents.classifier import ClassifierAgent

        agent = ClassifierAgent()
        agent.client = make_mock_client(MOCK_CLASSIFICATION_JSON)

        result = await agent.analyze(sample_mining_text)
        assert result["category"] == "safety_protocol"

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_confidence_is_float_in_range(self, sample_mining_text):
        """Confidence value is a float between 0 and 1."""
        from app.agents.classifier import ClassifierAgent

        agent = ClassifierAgent()
        agent.client = make_mock_client(MOCK_CLASSIFICATION_JSON)

        result = await agent.analyze(sample_mining_text)
        assert isinstance(result["confidence"], float)
        assert 0.0 <= result["confidence"] <= 1.0

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_unknown_category_maps_to_other(self, sample_mining_text):
        """Unknown category value falls back to 'other'."""
        from app.agents.classifier import ClassifierAgent

        agent = ClassifierAgent()
        agent.client = make_mock_client(
            json.dumps(
                {
                    "category": "completely_unknown_xyz",
                    "confidence": 0.5,
                    "reasoning": "unknown",
                }
            )
        )

        result = await agent.analyze(sample_mining_text)
        assert result["category"] == "other"

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_handles_empty_response(self, sample_mining_text):
        """Empty JSON response returns 'other' without crashing."""
        from app.agents.classifier import ClassifierAgent

        agent = ClassifierAgent()
        agent.client = make_mock_client("{}")

        result = await agent.analyze(sample_mining_text)
        assert "category" in result
        assert result["confidence"] == 0.5  # default


class TestSafetyAnalyzerAgent:
    """Tests for SafetyAnalyzerAgent."""

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_returns_expected_keys(self, sample_mining_text):
        """SafetyAnalyzerAgent.analyze() returns dict with required keys."""
        from app.agents.safety_analyzer import SafetyAnalyzerAgent

        agent = SafetyAnalyzerAgent()
        agent.client = make_mock_client(MOCK_SAFETY_JSON)

        result = await agent.analyze(
            sample_mining_text, {"category": "safety_protocol"}
        )

        for key in (
            "score",
            "status",
            "hazards",
            "recommendations",
            "compliance_details",
        ):
            assert key in result

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_score_is_numeric(self, sample_mining_text):
        """Safety score must be a numeric type."""
        from app.agents.safety_analyzer import SafetyAnalyzerAgent

        agent = SafetyAnalyzerAgent()
        agent.client = make_mock_client(MOCK_SAFETY_JSON)

        result = await agent.analyze(sample_mining_text)
        assert isinstance(result["score"], (int, float))

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_hazards_is_list(self, sample_mining_text):
        """Hazards must be a list."""
        from app.agents.safety_analyzer import SafetyAnalyzerAgent

        agent = SafetyAnalyzerAgent()
        agent.client = make_mock_client(MOCK_SAFETY_JSON)

        result = await agent.analyze(sample_mining_text)
        assert isinstance(result["hazards"], list)


class TestEntityExtractorAgent:
    """Tests for EntityExtractorAgent."""

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_returns_all_entity_types(self, sample_mining_text):
        """EntityExtractorAgent returns all 6 entity categories."""
        from app.agents.entity_extractor import EntityExtractorAgent

        agent = EntityExtractorAgent()
        agent.client = make_mock_client(MOCK_ENTITIES_JSON)

        result = await agent.analyze(sample_mining_text)

        for key in (
            "equipment",
            "chemicals",
            "locations",
            "personnel",
            "dates",
            "regulations",
        ):
            assert key in result
            assert isinstance(result[key], list)

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_entity_count_is_correct(self, sample_mining_text):
        """entity_count equals sum of all entity lists."""
        from app.agents.entity_extractor import EntityExtractorAgent

        agent = EntityExtractorAgent()
        agent.client = make_mock_client(MOCK_ENTITIES_JSON)

        result = await agent.analyze(sample_mining_text)
        expected = sum(
            len(result[k])
            for k in (
                "equipment",
                "chemicals",
                "locations",
                "personnel",
                "dates",
                "regulations",
            )
        )
        assert result["entity_count"] == expected


class TestSummarizerAgent:
    """Tests for SummarizerAgent."""

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_analyze_returns_expected_keys(self, sample_mining_text):
        """SummarizerAgent returns summary, key_points, action_items."""
        from app.agents.summarizer import SummarizerAgent

        agent = SummarizerAgent()
        agent.model = make_mock_model(MOCK_SUMMARY_JSON)

        result = await agent.analyze(sample_mining_text)

        assert "summary" in result
        assert "key_points" in result
        assert isinstance(result["summary"], str)
        assert len(result["summary"]) > 0

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_key_points_is_list(self, sample_mining_text):
        """key_points must be a list of strings."""
        from app.agents.summarizer import SummarizerAgent

        agent = SummarizerAgent()
        agent.model = make_mock_model(MOCK_SUMMARY_JSON)

        result = await agent.analyze(sample_mining_text)
        assert isinstance(result["key_points"], list)


class TestAgentOrchestrator:
    """Tests for AgentOrchestrator."""

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_orchestrator_returns_all_sections(self, sample_mining_text):
        """Orchestrator returns classification, safety, entities, summary, metadata."""
        from app.agents.orchestrator import AgentOrchestrator

        orchestrator = AgentOrchestrator()
        orchestrator.classifier.client = make_mock_client(MOCK_CLASSIFICATION_JSON)
        orchestrator.safety_analyzer.client = make_mock_client(MOCK_SAFETY_JSON)
        orchestrator.entity_extractor.client = make_mock_client(MOCK_ENTITIES_JSON)
        orchestrator.summarizer.model = make_mock_model(MOCK_SUMMARY_JSON)

        result = await orchestrator.analyze_document(sample_mining_text)

        for section in ("classification", "safety", "entities", "summary", "metadata"):
            assert section in result

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_orchestrator_includes_processing_time(self, sample_mining_text):
        """Metadata includes processing_time_ms > 0."""
        from app.agents.orchestrator import AgentOrchestrator

        orchestrator = AgentOrchestrator()
        orchestrator.classifier.client = make_mock_client(MOCK_CLASSIFICATION_JSON)
        orchestrator.safety_analyzer.client = make_mock_client(MOCK_SAFETY_JSON)
        orchestrator.entity_extractor.client = make_mock_client(MOCK_ENTITIES_JSON)
        orchestrator.summarizer.model = make_mock_model(MOCK_SUMMARY_JSON)

        result = await orchestrator.analyze_document(sample_mining_text)

        assert "processing_time_ms" in result["metadata"]
        assert isinstance(result["metadata"]["processing_time_ms"], int)
        assert result["metadata"]["processing_time_ms"] >= 0

    @pytest.mark.unit
    @pytest.mark.asyncio
    async def test_orchestrator_handles_agent_failure_gracefully(
        self, sample_mining_text
    ):
        """If one agent fails, orchestrator still returns results from other agents."""
        from app.agents.orchestrator import AgentOrchestrator

        orchestrator = AgentOrchestrator()
        orchestrator.classifier.client = make_mock_client(MOCK_CLASSIFICATION_JSON)
        orchestrator.safety_analyzer.client = make_mock_client(MOCK_SAFETY_JSON)
        orchestrator.entity_extractor.client = make_mock_client(MOCK_ENTITIES_JSON)

        # Make summarizer fail
        fail_model = MagicMock()
        fail_model.generate_content.side_effect = Exception("API rate limit exceeded")
        orchestrator.summarizer.model = fail_model

        result = await orchestrator.analyze_document(sample_mining_text)

        # Should still return — summary will have "error" key
        assert "summary" in result
        assert "error" in result["summary"] or "summary" in result["summary"]