test_procurement.py

"""
Unit and Integration Tests for Multi-Agent Procurement System

Run with: pytest test_procurement.py -v

Tests cover:
- Individual agent node behavior
- State transitions and routing
- Budget validation logic
- Graph compilation
- Error handling
"""

import pytest
from typing import Optional

from langgraph.graph import StateGraph, START, END

from procurement_system import (
    ProcurementState,
    ProcurementAgents,
    build_procurement_graph,
    ProcurementWorkflowExecutor,
    create_llm,
    mock_vendor_search,
    validate_budget,
    approval_gate,
    _approval_router,
)


# =============================================================================
# FIXTURES
# =============================================================================

@pytest.fixture
def llm():
    """Initialize language model for tests (Groq)."""
    return create_llm()


@pytest.fixture
def agents(llm):
    """Initialize agents for testing."""
    return ProcurementAgents(llm)


@pytest.fixture
def compiled_graph(llm):
    """Compile procurement graph."""
    graph, memory, agents = build_procurement_graph(llm)
    return graph, memory, agents


@pytest.fixture
def sample_state():
    """Create a sample procurement state."""
    return ProcurementState(
        procurement_request="Enterprise cloud infrastructure",
        vendor_options=[],
        selected_vendor={},
        budget_limit=5000.0,
        analysis_approved=False,
        human_approved=False,
        contract_draft="",
        logs=[]
    )


@pytest.fixture
def state_with_vendors(sample_state):
    """Create state with vendor options populated."""
    state = sample_state.copy()
    state["vendor_options"] = [
        {
            "vendor_id": "VENDOR_001",
            "name": "CloudTech Solutions",
            "service_type": "Cloud Infrastructure",
            "price_per_month": 3500.0,
            "capabilities": ["Auto-scaling", "99.99% Uptime"],
            "reputation_score": 9.2,
            "contract_terms": "12-month minimum"
        },
        {
            "vendor_id": "VENDOR_002",
            "name": "InfraSpeed Inc",
            "service_type": "Cloud Infrastructure",
            "price_per_month": 2800.0,
            "capabilities": ["Auto-scaling", "99.9% Uptime"],
            "reputation_score": 8.5,
            "contract_terms": "6-month minimum"
        }
    ]
    return state


# =============================================================================
# TOOL TESTS
# =============================================================================

class TestTools:
    """Test LangChain @tool functions."""

    def test_mock_vendor_search_cloud(self):
        """Test vendor search for cloud infrastructure."""
        result = mock_vendor_search("cloud infrastructure")

        assert isinstance(result, list)
        assert len(result) > 0
        assert all("vendor_id" in v for v in result)
        assert all("price_per_month" in v for v in result)

    def test_mock_vendor_search_software(self):
        """Test vendor search for software licensing."""
        result = mock_vendor_search("software licensing")

        assert isinstance(result, list)
        assert len(result) > 0

    def test_mock_vendor_search_returns_complete_data(self):
        """Verify vendor objects contain all required fields."""
        result = mock_vendor_search("cloud")
        vendor = result[0]

        required_fields = [
            "vendor_id", "name", "service_type", "price_per_month",
            "capabilities", "reputation_score", "contract_terms"
        ]
        assert all(field in vendor for field in required_fields)

    def test_validate_budget_within_limit(self):
        """Test budget validation when within limit."""
        result = validate_budget(vendor_price=3000.0, budget_limit=5000.0)

        assert result["within_budget"] == True
        assert result["vendor_price"] == 3000.0
        assert result["budget_limit"] == 5000.0
        assert result["variance"] == 2000.0
        assert result["status"] == "APPROVED"

    def test_validate_budget_exceeds_limit(self):
        """Test budget validation when exceeding limit."""
        result = validate_budget(vendor_price=6000.0, budget_limit=5000.0)

        assert result["within_budget"] == False
        assert result["variance"] == -1000.0
        assert result["variance_percentage"] == -20.0
        assert result["status"] == "REJECTED"

    def test_validate_budget_exact_limit(self):
        """Test budget validation at exact limit."""
        result = validate_budget(vendor_price=5000.0, budget_limit=5000.0)

        assert result["within_budget"] == True
        assert result["variance"] == 0.0
        assert result["variance_percentage"] == 0.0


# =============================================================================
# AGENT NODE TESTS
# =============================================================================

class TestAgentNodes:
    """Test individual agent node behavior."""

    def test_research_node_populates_vendors(self, agents, sample_state):
        """Test that research node finds vendors."""
        result = agents.research_node(sample_state)

        # Result should be a Command
        assert hasattr(result, 'update')
        assert hasattr(result, 'goto')

        # Should route to analysis
        assert result.goto == "analysis_node"

        # Should populate vendor_options
        assert "vendor_options" in result.update
        assert len(result.update["vendor_options"]) > 0

    def test_research_node_creates_log_entry(self, agents, sample_state):
        """Test that research node creates audit log."""
        result = agents.research_node(sample_state)

        assert "logs" in result.update
        assert len(result.update["logs"]) > 0
        assert "ResearchNode" in result.update["logs"][0]

    def test_analysis_node_within_budget(self, agents, state_with_vendors):
        """Test analysis node with budget-acceptable vendor."""
        result = agents.analysis_node(state_with_vendors)

        assert hasattr(result, 'update')
        assert result.goto == "approval_gate"
        assert result.update["analysis_approved"] == True
        assert "selected_vendor" in result.update

    def test_analysis_node_exceeds_budget(self, agents, state_with_vendors):
        """Test analysis node with budget-exceeded vendor."""
        # Reduce budget to force failure
        state_with_vendors["budget_limit"] = 2000.0

        result = agents.analysis_node(state_with_vendors)

        assert result.goto == END
        assert result.update["analysis_approved"] == False

    def test_legal_node_generates_contract(self, agents, sample_state):
        """Test legal node generates contract draft."""
        sample_state["selected_vendor"] = {
            "name": "CloudTech Solutions",
            "service_type": "Cloud Infrastructure",
            "price_per_month": 5000.0,
            "capabilities": ["Auto-scaling", "99.99% Uptime SLA"],
            "contract_terms": "12-month minimum"
        }
        sample_state["analysis_approved"] = True

        result = agents.legal_node(sample_state)

        assert "contract_draft" in result.update
        assert len(result.update["contract_draft"]) > 0
        assert "PURCHASE" in result.update["contract_draft"].upper() or "CONTRACT" in result.update["contract_draft"].upper()
        assert result.goto == END


# =============================================================================
# ROUTING AND CONDITIONAL LOGIC TESTS
# =============================================================================

class TestRouting:
    """Test graph routing and conditional logic."""

    def test_approval_gate_returns_dict_with_logs(self, sample_state):
        """Test approval gate returns dict with a logs entry."""
        sample_state["selected_vendor"] = {"name": "TestVendor"}
        result = approval_gate(sample_state)
        assert isinstance(result, dict)
        assert "logs" in result

    def test_approval_router_routes_to_legal_when_approved(self, sample_state):
        """Test _approval_router returns 'legal_node' when human_approved=True."""
        sample_state["human_approved"] = True
        assert _approval_router(sample_state) == "legal_node"

    def test_approval_router_routes_to_end_when_rejected(self, sample_state):
        """Test _approval_router returns END when human_approved=False."""
        sample_state["human_approved"] = False
        assert _approval_router(sample_state) == END


# =============================================================================
# GRAPH COMPILATION TESTS
# =============================================================================

class TestGraphCompilation:
    """Test graph building and compilation."""

    def test_graph_compiles_successfully(self, compiled_graph):
        """Test that graph compiles without errors."""
        graph, memory, agents = compiled_graph

        assert graph is not None
        assert memory is not None
        assert agents is not None

    def test_compiled_graph_has_required_nodes(self, compiled_graph):
        """Test that compiled graph has all required nodes."""
        graph, _, _ = compiled_graph

        # Get graph schema
        schema = graph.get_schema()

        # Verify graph is callable
        assert callable(graph.invoke) or callable(graph.stream)

    def test_graph_has_interrupt_before_legal(self, compiled_graph):
        """Test that graph is configured to interrupt before legal_node."""
        graph, _, _ = compiled_graph

        # This is implicit in the compilation - we can test by checking
        # that the graph actually pauses at the right point
        # (This is verified through integration tests)
        assert graph is not None


# =============================================================================
# INTEGRATION TESTS
# =============================================================================

class TestIntegration:
    """Integration tests of complete workflow."""

    def test_workflow_execution_happy_path(self, compiled_graph):
        """Test complete workflow from start to approval point."""
        graph, memory, _ = compiled_graph
        executor = ProcurementWorkflowExecutor(graph, memory)

        state, _ = executor.start_workflow(
            procurement_request="Cloud infrastructure",
            budget_limit=5500.0,
        )

        assert len(state["vendor_options"]) > 0
        assert state["selected_vendor"]["name"] is not None
        assert state["analysis_approved"] == True
        assert state["human_approved"] == False
        assert len(state["logs"]) > 0

    def test_workflow_budget_rejection(self, compiled_graph):
        """Test workflow rejection when budget is exceeded."""
        graph, memory, _ = compiled_graph
        executor = ProcurementWorkflowExecutor(graph, memory)

        state, _ = executor.start_workflow(
            procurement_request="Cloud infrastructure",
            budget_limit=2000.0,
        )

        assert state["analysis_approved"] == False
        assert state["contract_draft"] == ""

    def test_workflow_approval_and_resumption(self, compiled_graph):
        """Test workflow resumption after human approval."""
        graph, memory, _ = compiled_graph
        executor = ProcurementWorkflowExecutor(graph, memory)

        state, _ = executor.start_workflow(
            procurement_request="Cloud infrastructure",
            budget_limit=5500.0,
        )
        assert state["human_approved"] == False

        final_state, _ = executor.approve_vendor(approval=True)
        assert final_state["human_approved"] == True
        assert len(final_state["contract_draft"]) > 0


# =============================================================================
# STATE MANAGEMENT TESTS
# =============================================================================

class TestStateManagement:
    """Test state persistence and management."""

    def test_state_checkpoint_created(self, compiled_graph):
        """Test that state is checkpointed at interruption."""
        graph, memory, _ = compiled_graph
        executor = ProcurementWorkflowExecutor(graph, memory)

        executor.start_workflow(
            procurement_request="Cloud infrastructure",
            budget_limit=5500.0,
        )
        config      = {"configurable": {"thread_id": executor.thread_id}}
        saved_state = graph.get_state(config)
        assert saved_state is not None
        assert saved_state.values is not None

    def test_state_recovery_after_interruption(self, compiled_graph):
        """Test state can be retrieved after interruption."""
        graph, memory, _ = compiled_graph
        executor = ProcurementWorkflowExecutor(graph, memory)

        state, _ = executor.start_workflow(
            procurement_request="Cloud infrastructure",
            budget_limit=5500.0,
        )
        current = executor.get_state()
        assert current["budget_limit"] == state["budget_limit"]
        assert current["selected_vendor"]["name"] == state["selected_vendor"]["name"]


# =============================================================================
# ERROR HANDLING TESTS
# =============================================================================

class TestErrorHandling:
    """Test error handling and edge cases."""

    def test_empty_vendor_list_handling(self, agents, sample_state):
        """Test handling of empty vendor search results."""
        # This shouldn't normally happen with mock_vendor_search
        # but real implementations might fail
        sample_state["vendor_options"] = []

        # Should handle gracefully (implementation dependent)
        # Current implementation would route to analysis_node
        # where it would select first vendor (if available)

    def test_missing_vendor_fields(self, agents, sample_state):
        """Test handling of incomplete vendor data."""
        sample_state["vendor_options"] = [
            {
                "vendor_id": "VENDOR_001",
                # Missing required fields
                "name": "Incomplete Vendor"
            }
        ]

        # Should handle gracefully with defaults
        result = agents.analysis_node(sample_state)

        # Should still route (with defaults for missing fields)
        assert hasattr(result, 'goto')

    def test_invalid_budget_values(self):
        """Test budget validation with edge case values."""
        # Zero budget
        result = validate_budget(vendor_price=1000.0, budget_limit=0.0)
        assert result["within_budget"] == False

        # Negative values
        result = validate_budget(vendor_price=-100.0, budget_limit=1000.0)
        assert result["within_budget"] == True  # Negative price within any budget

        # Very large values
        result = validate_budget(vendor_price=999999.99, budget_limit=1000000.0)
        assert result["within_budget"] == True


# =============================================================================
# PERFORMANCE TESTS
# =============================================================================

class TestPerformance:
    """Test performance characteristics."""

    def test_vendor_search_latency(self):
        """Test vendor search completes quickly."""
        import time

        start = time.time()
        result = mock_vendor_search("cloud")
        elapsed = time.time() - start

        # Should be very fast (mock)
        assert elapsed < 1.0
        assert len(result) > 0

    def test_budget_validation_latency(self):
        """Test budget validation completes quickly."""
        import time

        start = time.time()
        result = validate_budget(5000.0, 5500.0)
        elapsed = time.time() - start

        # Should be instant
        assert elapsed < 0.1
        assert result["within_budget"] == True


# =============================================================================
# CONFIGURATION TESTS
# =============================================================================

class TestConfiguration:
    """Test system configuration and environment."""

    def test_langsmith_config_detection(self):
        """Test that LangSmith config can be detected."""
        import os

        # Should detect without error even if not configured
        api_key = os.getenv("LANGSMITH_API_KEY", "")
        project = os.getenv("LANGSMITH_PROJECT", "")

        # Test passes if no exceptions
        assert isinstance(api_key, str)
        assert isinstance(project, str)

    def test_groq_llm_available(self, llm):
        """Test that Groq LLM is configured and invocable."""
        assert llm is not None
        assert hasattr(llm, 'invoke')


# =============================================================================
# MAIN EXECUTION
# =============================================================================

if __name__ == "__main__":
    # Run tests with verbose output
    pytest.main([__file__, "-v", "--tb=short"])