Chroma MCP Server

""" Tests for the ThinkingSession class. """ import pytest from mcp import ClientSession from unittest.mock import MagicMock, patch # Mock the dependencies # Assume McpClient needs to be mocked if ThinkingSession instantiates it # If a client is passed in, we might mock that specific instance instead # Patching the class globally for simplicity here # with patch("chroma_mcp_thinking.thinking_session.ChromaMcpClient"): with patch("chroma_mcp_thinking.thinking_session.ClientSession") as MockMcpClient: from chroma_mcp_thinking.thinking_session import ThinkingSession @pytest.fixture def mock_client(): """Create a mock ChromaMcpClient.""" client = MagicMock() client.mcp_chroma_dev_chroma_sequential_thinking.return_value = {"session_id": "test-session-id"} client.mcp_chroma_dev_chroma_find_similar_thoughts.return_value = { "similar_thoughts": [ { "metadata": {"session_id": "test-session-id", "thought_number": 1}, "document": "Test thought", "distance": 0.5, } ] } client.mcp_chroma_dev_chroma_find_similar_sessions.return_value = { "similar_sessions": [ {"metadata": {"session_id": "test-session-id"}, "document": "First thought", "distance": 0.5} ] } client.mcp_chroma_dev_chroma_get_session_summary.return_value = { "thoughts": [{"metadata": {"thought_number": 1, "session_id": "test-session-id"}, "document": "Test thought"}] } return client def test_thinking_session_init(): """Test ThinkingSession initialization.""" # Test with default arguments # Patch ClientSession as used within ThinkingSession's __init__ with patch("chroma_mcp_thinking.thinking_session.ClientSession") as MockClientSession: mock_instance = MockClientSession.return_value mock_instance.mcp_chroma_dev_chroma_sequential_thinking = MagicMock() mock_instance.mcp_chroma_dev_chroma_find_similar_thoughts = MagicMock() mock_instance.mcp_chroma_dev_chroma_get_session_summary = MagicMock() mock_instance.mcp_chroma_dev_chroma_find_similar_sessions = MagicMock() session = ThinkingSession() # Should now use the mocked ClientSession assert isinstance(session.client, MagicMock) # Check it used the mocked instance assert session.session_id is not None MockClientSession.assert_called_once() # Verify ClientSession was instantiated # Test with provided session ID test_session_id = "test-session-123" with patch("chroma_mcp_thinking.thinking_session.ClientSession") as MockClientSession: session = ThinkingSession(session_id=test_session_id) assert session.session_id == test_session_id MockClientSession.assert_called_once() # Test with provided client instance mock_provided_client = MagicMock(spec=ClientSession) session = ThinkingSession(client=mock_provided_client, session_id="provided-client-session") assert session.client is mock_provided_client assert session.session_id == "provided-client-session" def test_record_thought(mock_client): """Test recording a thought.""" session = ThinkingSession(client=mock_client, session_id="test-session-id") result = session.record_thought( thought="Test thought", thought_number=1, total_thoughts=3, next_thought_needed=True ) # Client method should be called with correct arguments mock_client.mcp_chroma_dev_chroma_sequential_thinking.assert_called_once() call_args = mock_client.mcp_chroma_dev_chroma_sequential_thinking.call_args[1] assert call_args["thought"] == "Test thought" assert call_args["thought_number"] == 1 assert call_args["total_thoughts"] == 3 assert call_args["session_id"] == "test-session-id" assert call_args["next_thought_needed"] is True # Result should be returned from client method assert result == {"session_id": "test-session-id"} def test_find_similar_thoughts(mock_client): """Test finding similar thoughts.""" session = ThinkingSession(client=mock_client, session_id="test-session-id") results = session.find_similar_thoughts(query="test query", n_results=5, threshold=0.7, include_branches=True) # Client method should be called with correct arguments mock_client.mcp_chroma_dev_chroma_find_similar_thoughts.assert_called_once() call_args = mock_client.mcp_chroma_dev_chroma_find_similar_thoughts.call_args[1] assert call_args["query"] == "test query" assert call_args["session_id"] == "test-session-id" assert call_args["n_results"] == 5 assert call_args["threshold"] == 0.7 assert call_args["include_branches"] is True # Results should be extracted from client response assert len(results) == 1 assert results[0]["document"] == "Test thought" assert results[0]["metadata"]["session_id"] == "test-session-id" def test_get_session_summary(mock_client): """Test getting session summary.""" session = ThinkingSession(client=mock_client, session_id="test-session-id") summary = session.get_session_summary(include_branches=True) # Client method should be called with correct arguments mock_client.mcp_chroma_dev_chroma_get_session_summary.assert_called_once() call_args = mock_client.mcp_chroma_dev_chroma_get_session_summary.call_args[1] assert call_args["session_id"] == "test-session-id" assert call_args["include_branches"] is True # Summary should be returned from client method assert "thoughts" in summary assert len(summary["thoughts"]) == 1 assert summary["thoughts"][0]["document"] == "Test thought" def test_find_similar_sessions(mock_client): """Test finding similar sessions.""" results = ThinkingSession.find_similar_sessions(query="test query", client=mock_client, n_results=5, threshold=0.7) # Client method should be called with correct arguments mock_client.mcp_chroma_dev_chroma_find_similar_sessions.assert_called_once() call_args = mock_client.mcp_chroma_dev_chroma_find_similar_sessions.call_args[1] assert call_args["query"] == "test query" assert call_args["n_results"] == 5 assert call_args["threshold"] == 0.7 # Results should be extracted from client response assert len(results) == 1 assert results[0]["document"] == "First thought" assert results[0]["metadata"]["session_id"] == "test-session-id"