mcp-forge/tests/server/tools/test_document_state.py

"""Tests for Document State Tool."""

import pytest
from unittest.mock import Mock, AsyncMock
from mcp.types import Tool, TextContent
import json

from mcp_forge.server.tools.document_state import DocumentStateTool


@pytest.fixture
def mock_session_manager():
    """Mock session manager."""
    manager = Mock()
    # Mock get_session_state to return a session with state
    mock_session = Mock()
    mock_session.state = Mock()
    mock_session.state.to_dict = Mock(return_value={
        "variables": {"x": 1, "y": 2},
        "documented_variables": {},
        "note": None
    })
    manager.get_session_state = Mock(return_value=mock_session)
    manager.document_variables = AsyncMock(return_value={"success": True, "documented_count": 2})
    manager.session_exists = Mock(return_value=True)
    return manager


@pytest.fixture
def document_state_tool(mock_session_manager):
    """Create DocumentStateTool instance with mocked dependencies."""
    return DocumentStateTool(session_manager=mock_session_manager)


@pytest.mark.asyncio
async def test_get_tool_definition(document_state_tool):
    """Test that tool definition matches MCP spec."""
    definition = document_state_tool.get_tool_definition()
    
    assert isinstance(definition, Tool)
    assert definition.name == "document_state"
    assert definition.description is not None
    assert "document" in definition.description.lower()
    
    # Verify required schema properties
    schema = definition.inputSchema
    assert schema["type"] == "object"
    assert "session_id" in schema["properties"]
    assert "variables" in schema["properties"]
    assert "note" in schema["properties"]
    assert "clear" in schema["properties"]
    assert "session_id" in schema["required"]
    assert "variables" in schema["required"]


@pytest.mark.asyncio
async def test_document_variables(
    document_state_tool,
    mock_session_manager
):
    """Test documenting variables in a session."""
    arguments = {
        "session_id": "test-session",
        "variables": {
            "df": "Customer data with 1000 rows",
            "model": "Trained RandomForest classifier"
        }
    }
    
    result = await document_state_tool.execute(arguments)
    
    assert isinstance(result, list)
    assert len(result) == 1
    assert isinstance(result[0], TextContent)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    assert response["documented_count"] == 2
    assert response["session_id"] == "test-session"
    
    # Verify session manager was called
    mock_session_manager.document_variables.assert_called_once()


@pytest.mark.asyncio
async def test_document_with_note(
    document_state_tool,
    mock_session_manager
):
    """Test documenting with a note."""
    arguments = {
        "session_id": "test-session",
        "variables": {
            "result": "Final analysis output"
        },
        "note": "Analysis complete, ready for reporting"
    }
    
    result = await document_state_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    
    # Verify note was passed
    call_args = mock_session_manager.document_variables.call_args
    assert call_args is not None
    assert call_args.kwargs["note"] == "Analysis complete, ready for reporting"


@pytest.mark.asyncio
async def test_clear_existing_documentation(
    document_state_tool,
    mock_session_manager
):
    """Test clearing existing documentation."""
    arguments = {
        "session_id": "test-session",
        "variables": {
            "new_var": "New variable"
        },
        "clear": True
    }
    
    result = await document_state_tool.execute(arguments)
    
    # Parse JSON response
    response = json.loads(result[0].text)
    assert response["success"] is True
    
    # Verify clear flag was passed
    call_args = mock_session_manager.document_variables.call_args
    assert call_args is not None
    assert call_args.kwargs["clear"] is True


@pytest.mark.asyncio
async def test_session_validation(
    document_state_tool,
    mock_session_manager
):
    """Test that validation fails for non-existent session."""
    mock_session_manager.session_exists = Mock(return_value=False)
    
    arguments = {
        "session_id": "nonexistent",
        "variables": {
            "x": "test"
        }
    }
    
    with pytest.raises(ValueError, match="Session"):
        await document_state_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_missing_session_id(document_state_tool):
    """Test that validation fails when session_id is missing."""
    arguments = {
        "variables": {"x": "test"}
    }
    
    with pytest.raises((ValueError, KeyError)):
        await document_state_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_missing_variables(document_state_tool):
    """Test that validation fails when variables is missing."""
    arguments = {
        "session_id": "test-session"
    }
    
    with pytest.raises((ValueError, KeyError)):
        await document_state_tool.execute(arguments)


@pytest.mark.asyncio
async def test_validate_arguments_invalid_variables_type(document_state_tool):
    """Test that validation fails when variables is not a dict."""
    arguments = {
        "session_id": "test-session",
        "variables": ["not", "a", "dict"]
    }
    
    with pytest.raises(ValueError, match="variables"):
        await document_state_tool.execute(arguments)