initial commit after one day coding agent session

tests/server/tools/test_document_state.py

@@ -0,0 +1,188 @@
+"""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)