test_thinking_tools.py•39.8 kB
"""Tests for thinking tools."""
import pytest
import uuid
import time
import json
from contextlib import contextmanager
from typing import Dict, Any, List, Optional
from datetime import datetime
from unittest.mock import patch, MagicMock, ANY
from mcp import types as mcp_types
from mcp.shared.exceptions import McpError
from mcp.types import INVALID_PARAMS, INTERNAL_ERROR, ErrorData
from src.chroma_mcp.utils.errors import ValidationError
from src.chroma_mcp.types import ChromaClientConfig
from src.chroma_mcp.tools.thinking_tools import (
ThoughtMetadata, # Import if needed for checks
THOUGHTS_COLLECTION, # Import constants
SESSIONS_COLLECTION, # FIX: Import this constant
DEFAULT_SIMILARITY_THRESHOLD,
_sequential_thinking_impl,
_find_similar_thoughts_impl,
_get_session_summary_impl,
_find_similar_sessions_impl,
SequentialThinkingInput,
FindSimilarThoughtsInput,
GetSessionSummaryInput,
FindSimilarSessionsInput,
SequentialThinkingWithCustomDataInput,
_sequential_thinking_with_custom_data_impl,
)
# --- Helper Functions (Copied from test_collection_tools.py) ---
def assert_successful_json_list_result(
result: List[mcp_types.TextContent], expected_data: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""Asserts the result is a list containing one TextContent with valid JSON."""
assert isinstance(result, list)
assert len(result) == 1
assert isinstance(result[0], mcp_types.TextContent)
assert result[0].type == "text"
try:
result_data = json.loads(result[0].text)
except json.JSONDecodeError:
pytest.fail(f"Failed to parse JSON content: {result[0].text}")
assert isinstance(result_data, dict)
if expected_data is not None:
assert result_data == expected_data
return result_data
# Context manager to assert McpError is raised
@contextmanager
def assert_raises_mcp_error(expected_message: str):
"""Context manager to check if a specific McpError is raised."""
try:
yield
except McpError as e:
# Check both e.error_data (if exists) and e.args[0]
message = ""
if hasattr(e, "error_data") and hasattr(e.error_data, "message"):
message = str(e.error_data.message)
elif e.args and isinstance(e.args[0], ErrorData) and hasattr(e.args[0], "message"):
message = str(e.args[0].message)
else:
message = str(e) # Fallback to the exception string itself
assert (
expected_message in message
), f"Expected error message containing '{expected_message}' but got '{message}'"
return
except Exception as e:
pytest.fail(f"Expected McpError but got {type(e).__name__}: {e}")
pytest.fail("Expected McpError but no exception was raised.")
# --- End Helper Functions ---
@pytest.fixture
def mock_chroma_client_thinking():
"""Provides a mocked Chroma client, thoughts collection, and sessions collection.
Also ensures that get_server_config within thinking_tools returns a valid config.
"""
# Patch get_chroma_client within the thinking_tools module
with (
patch("src.chroma_mcp.tools.thinking_tools.get_chroma_client") as mock_get_client,
patch("src.chroma_mcp.tools.thinking_tools.get_server_config") as mock_tt_get_server_config,
):
# Setup a default ChromaClientConfig for the tests
# This is what get_server_config() will return when called by the _impl functions
# within thinking_tools.py
test_config = ChromaClientConfig(client_type="ephemeral", embedding_function_name="default")
mock_tt_get_server_config.return_value = test_config
mock_client = MagicMock()
mock_thoughts_collection = MagicMock(name="thoughts_collection")
mock_sessions_collection = MagicMock(name="sessions_collection")
# Configure the mock client to return specific collections by name
def get_collection_side_effect(name, embedding_function=None):
if name == THOUGHTS_COLLECTION:
return mock_thoughts_collection
elif name == SESSIONS_COLLECTION:
return mock_sessions_collection
else:
# Raise error for unexpected collection names
raise ValueError(f"Mock Client: Collection {name} does not exist.")
mock_client.get_collection.side_effect = get_collection_side_effect
mock_client.get_or_create_collection.side_effect = get_collection_side_effect # Mock this too
mock_get_client.return_value = mock_client
# Also patch the embedding function if it's used directly
with patch("src.chroma_mcp.tools.thinking_tools.get_embedding_function") as mock_get_emb:
mock_get_emb.return_value = MagicMock(name="mock_embedding_function")
# Yield all mocks needed by tests
yield mock_client, mock_thoughts_collection, mock_sessions_collection
@pytest.mark.usefixtures("initialized_chroma_client")
class TestThinkingTools:
"""Test cases for thinking tools implementation functions."""
# --- _sequential_thinking_impl Tests ---
@pytest.mark.asyncio # Mark as async
async def test_sequential_thinking_new_session_success(self, mock_chroma_client_thinking):
"""Test recording the first thought in a new session."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
input_data = SequentialThinkingInput(
thought="Initial thought",
thought_number=1,
total_thoughts=3,
)
# ACT
result = await _sequential_thinking_impl(input_data)
# ASSERT using MODIFIED helper
result_data = assert_successful_json_list_result(result)
# Assertions on mocks
mock_client.get_or_create_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
mock_collection.add.assert_called_once()
call_args = mock_collection.add.call_args
assert call_args.kwargs["documents"] == ["Initial thought"]
assert len(call_args.kwargs["ids"]) == 1
thought_id = call_args.kwargs["ids"][0]
assert thought_id.startswith("thought_")
assert call_args.kwargs["metadatas"][0]["thought_number"] == 1
assert call_args.kwargs["metadatas"][0]["total_thoughts"] == 3
assert "session_id" in call_args.kwargs["metadatas"][0]
session_id = call_args.kwargs["metadatas"][0]["session_id"]
# Assertions on result data keys and values
assert result_data.get("thought_id") == thought_id
assert result_data.get("session_id") == session_id
assert result_data.get("previous_thoughts_count") == 0 # Check count
@pytest.mark.asyncio # Mark as async
async def test_sequential_thinking_existing_session(self, mock_chroma_client_thinking):
"""Test recording subsequent thoughts in an existing session."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
session_id = "existing_session"
# Mock get_collection returning empty results for previous thoughts check
# We need to mock `get` for the previous thoughts check
mock_collection.get.return_value = {"ids": [], "metadatas": [], "documents": []}
input_data = SequentialThinkingInput(
thought="Second thought",
thought_number=2,
total_thoughts=3,
session_id=session_id,
)
# ACT
result = await _sequential_thinking_impl(input_data)
# ASSERT using MODIFIED helper
result_data = assert_successful_json_list_result(result)
assert result_data["session_id"] == session_id
assert result_data["thought_id"].startswith(f"thought_{session_id}_2")
# Assertions on mocks
# Use keyword argument 'name'
mock_client.get_or_create_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
mock_collection.add.assert_called_once()
# Assert the get call for previous thoughts
mock_collection.get.assert_called_once()
call_args, call_kwargs = mock_collection.get.call_args
assert call_kwargs.get("where") == {"session_id": session_id}
@pytest.mark.asyncio # Mark as async
async def test_sequential_thinking_new_branch(self, mock_chroma_client_thinking):
"""Test starting a new branch from a specific thought."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
session_id = "branch_session_456"
branch_id = "feature_branch_abc"
branch_from = 2
input_data = SequentialThinkingInput(
thought="First thought on the branch",
thought_number=3,
total_thoughts=5,
session_id=session_id,
branch_id=branch_id,
branch_from_thought=branch_from,
)
# Mock collection.get for previous thoughts (might return thoughts from main trunk)
mock_collection.get.return_value = {
"ids": [f"thought_{session_id}_1", f"thought_{session_id}_2"],
"documents": ["First idea", "Second idea"],
"metadatas": [
{
"session_id": session_id,
"thought_number": 1,
"total_thoughts": 3,
"timestamp": 12345,
"branch_id": None,
},
{
"session_id": session_id,
"thought_number": 2,
"total_thoughts": 3,
"timestamp": 12345,
"branch_id": branch_id,
},
],
}
# ACT
result = await _sequential_thinking_impl(input_data)
# ASSERT using MODIFIED helper
result_data = assert_successful_json_list_result(result)
# Assertions on mocks
mock_collection.get.assert_called_once()
get_call_args = mock_collection.get.call_args
# Updated expected_where to reflect the simplified get call
expected_where = {"session_id": session_id} # Fetch all session thoughts, filter in Python
assert get_call_args.kwargs["where"] == expected_where
mock_collection.add.assert_called_once()
add_call_args = mock_collection.add.call_args
assert add_call_args.kwargs["metadatas"][0]["session_id"] == session_id
assert add_call_args.kwargs["metadatas"][0]["branch_id"] == branch_id
added_thought_id = add_call_args.kwargs["ids"][0]
# Assertions on result data
assert result_data.get("session_id") == session_id
assert result_data.get("thought_id") == added_thought_id
assert result_data.get("previous_thoughts_count") == 2 # Check count
@pytest.mark.asyncio # Mark as async
async def test_sequential_thinking_invalid_custom_data_json(self, mock_chroma_client_thinking):
"""Test failure when custom_data is invalid JSON."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
input_data = SequentialThinkingWithCustomDataInput(
thought="Thought with bad custom data",
thought_number=1,
total_thoughts=1,
custom_data='{"key": "value", "unterminated}', # Invalid JSON
)
# ASSERT: Expect McpError to be raised now
with assert_raises_mcp_error("Invalid JSON format for custom_data string"):
await _sequential_thinking_with_custom_data_impl(input_data)
# Ensure add was not called
mock_collection.add.assert_not_called()
@pytest.mark.asyncio # Mark as async
async def test_sequential_thinking_custom_data_not_dict(self, mock_chroma_client_thinking):
"""Test failure when custom_data JSON decodes to non-dictionary."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
input_data = SequentialThinkingWithCustomDataInput(
thought="Thought with non-dict custom data",
thought_number=1,
total_thoughts=1,
custom_data="[1, 2, 3]", # Valid JSON, but not a dict
)
# ASSERT: Expect McpError to be raised now
with assert_raises_mcp_error("Custom data string must decode to a JSON object (dictionary)."):
await _sequential_thinking_with_custom_data_impl(input_data)
# Ensure add was not called
mock_collection.add.assert_not_called()
# --- _find_similar_thoughts_impl Tests ---
@pytest.mark.asyncio # Mark as async
async def test_find_similar_thoughts_success(self, mock_chroma_client_thinking):
"""Test finding similar thoughts successfully."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
query = "find me similar ideas"
threshold = DEFAULT_SIMILARITY_THRESHOLD # Use constant
# Mock query results
mock_collection.query.return_value = {
"ids": [["t1", "t2", "t3"]],
"documents": [["idea A", "idea B", "idea C"]],
"metadatas": [[{"session_id": "s1"}, {"session_id": "s2", "custom:tag": "X"}, {"session_id": "s1"}]],
"distances": [[0.1, 0.25, 0.4]], # Similarities: 0.9, 0.75, 0.6
}
# ACT - Use default n_results from model (which is 5)
input_model = FindSimilarThoughtsInput(query=query, threshold=threshold)
result = await _find_similar_thoughts_impl(input_model)
# ASSERT using MODIFIED helper
result_data = assert_successful_json_list_result(result)
# Assertions on mocks
mock_client.get_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
# Expect n_results=5 (Pydantic default)
mock_collection.query.assert_called_once_with(
query_texts=[query], n_results=5, where=None, include=["documents", "metadatas", "distances"]
)
# Assertions on result data
assert (
len(result_data.get("similar_thoughts", [])) == 3
) # t3 (similarity 0.6) is now included due to threshold change
assert result_data.get("total_found") == 3
assert result_data.get("threshold_used") == threshold
# Check content of results
thought1 = result_data["similar_thoughts"][0]
assert thought1["id"] == "t1"
assert thought1["similarity"] == 0.9
assert thought1["metadata"]["session_id"] == "s1"
# Check reconstructed custom data
thought2 = result_data["similar_thoughts"][1]
assert thought2["id"] == "t2"
assert thought2["metadata"]["session_id"] == "s2"
assert thought2["metadata"].get("custom_data") == {"tag": "X"}
@pytest.mark.asyncio # Mark as async
async def test_find_similar_thoughts_with_session_filter(self, mock_chroma_client_thinking):
"""Test finding similar thoughts filtered by session ID."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
session_id_to_find = "s1"
# Mock query results (same as above)
mock_collection.query.return_value = {
"ids": [["t1", "t2", "t3"]],
"documents": [["idea A", "idea B", "idea C"]],
"metadatas": [[{"session_id": "s1"}, {"session_id": "s2"}, {"session_id": "s1"}]],
"distances": [[0.1, 0.25, 0.4]],
}
# ACT
input_model = FindSimilarThoughtsInput(
query=session_id_to_find, session_id=session_id_to_find, n_results=3, include_branches=False, threshold=0.5
)
result = await _find_similar_thoughts_impl(input_model)
# ASSERT using MODIFIED helper
result_data = assert_successful_json_list_result(result)
# Check query was called with the where clause
mock_collection.query.assert_called_once_with(
query_texts=[session_id_to_find], n_results=3, where={"session_id": session_id_to_find}, include=ANY
)
# Assertions on result data (mock doesn't filter, so we expect all 3, but check metadata)
assert len(result_data.get("similar_thoughts", [])) == 3
assert result_data["similar_thoughts"][0].get("metadata", {}).get("session_id") == "s1"
# assert result_data["similar_thoughts"][1].get("metadata", {}).get("session_id") == "s2" # Mock doesn't filter
assert result_data["similar_thoughts"][2].get("metadata", {}).get("session_id") == "s1"
@pytest.mark.asyncio # Mark as async
async def test_find_similar_thoughts_collection_not_found(self, mock_chroma_client_thinking):
"""Test find_similar_thoughts when the collection does not exist."""
mock_client, _, _ = mock_chroma_client_thinking
# Mock get_collection to raise the specific ValueError
mock_client.get_collection.side_effect = ValueError(f"Collection {THOUGHTS_COLLECTION} does not exist.")
# ACT
input_model = FindSimilarThoughtsInput(query="any query")
result = await _find_similar_thoughts_impl(input_model)
# ASSERT: Should return success list with a message
result_data = assert_successful_json_list_result(result)
assert result_data.get("similar_thoughts") == []
assert result_data.get("total_found") == 0
mock_client.get_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
# Ensure query was not called if collection not found
# (Need the mock collection instance for this)
# Assuming the side effect prevents query from being called implicitly
# If side_effect doesn't prevent access, mock collection mock would be needed:
# mock_thoughts_collection.query.assert_not_called()
# --- _get_session_summary_impl Tests ---
@pytest.mark.asyncio # Mark as async
async def test_get_session_summary_success(self, mock_chroma_client_thinking):
"""Test getting a session summary successfully."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
session_id = "summary_session"
# Mock collection.get results
mock_collection.get.return_value = {
"ids": [f"thought_{session_id}_2", f"thought_{session_id}_1"], # Unordered
"documents": ["Thought two", "Thought one"],
"metadatas": [
{"session_id": session_id, "thought_number": 2, "custom:tag": "final"},
{"session_id": session_id, "thought_number": 1},
],
}
# ACT
input_model = GetSessionSummaryInput(session_id=session_id)
result = await _get_session_summary_impl(input_model)
# ASSERT using MODIFIED helper
result_data = assert_successful_json_list_result(result)
# Assertions on mocks
mock_client.get_or_create_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
mock_collection.get.assert_called_once_with(include=["documents", "metadatas"])
# Assertions on result data
assert result_data["session_id"] == session_id
assert result_data["total_thoughts_in_session"] == 2
assert len(result_data["session_thoughts"]) == 2
# Check order (should be sorted by thought_number)
assert result_data["session_thoughts"][0]["content"] == "Thought one"
assert result_data["session_thoughts"][1]["content"] == "Thought two"
# Check reconstructed custom data
assert result_data["session_thoughts"][1]["metadata"].get("custom_data") == {"tag": "final"}
@pytest.mark.asyncio # Mark as async
async def test_get_session_summary_no_thoughts(self, mock_chroma_client_thinking):
"""Test getting a summary for a session with no thoughts found."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
session_id = "empty_session"
# Mock collection.get returning empty results
mock_collection.get.return_value = {"ids": [], "documents": [], "metadatas": []}
# Expected result structure when no thoughts are found
expected_result_data = {"session_id": session_id, "session_thoughts": [], "total_thoughts_in_session": 0}
# ACT
input_model = GetSessionSummaryInput(session_id=session_id)
result = await _get_session_summary_impl(input_model)
# ASSERT using MODIFIED helper - compare with the expected structure
assert_successful_json_list_result(result, expected_data=expected_result_data)
# Assertions on mocks
mock_client.get_or_create_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
mock_collection.get.assert_called_once_with(include=["documents", "metadatas"])
@pytest.mark.asyncio # Mark as async
async def test_get_session_summary_collection_not_found(self, mock_chroma_client_thinking):
"""Test getting summary when the collection does not exist (but get_or_create succeeds)."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
session_id = "no_collection_session"
# Configure get_or_create_collection to return the mock collection
mock_client.get_or_create_collection.return_value = mock_collection
# Configure the collection's get method to return empty results
mock_collection.get.return_value = {"ids": [], "documents": [], "metadatas": []}
# ACT
input_model = GetSessionSummaryInput(session_id=session_id)
result = await _get_session_summary_impl(input_model)
# ASSERT: Expect success list with empty thoughts
result_data = assert_successful_json_list_result(result)
assert result_data.get("session_id") == session_id
assert result_data.get("session_thoughts") == []
assert result_data.get("total_thoughts_in_session") == 0
# Check message is NOT the 'not found' message
assert "not found" not in result_data.get("message", "")
# Assert mocks were called
mock_client.get_or_create_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
mock_collection.get.assert_called_once_with(include=["documents", "metadatas"])
@pytest.mark.asyncio
async def test_get_session_summary_unexpected_error(self, mock_chroma_client_thinking):
"""Test unexpected error during get session summary."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
error_message = "Cannot get thoughts"
mock_collection.get.side_effect = Exception(error_message)
# ACT & ASSERT: Expect McpError to be raised
input_model = GetSessionSummaryInput(session_id="some_session")
with assert_raises_mcp_error(
f"Tool Error: An unexpected error occurred while getting session summary for 'some_session'. Details: {error_message}"
):
await _get_session_summary_impl(input_model)
mock_collection.get.assert_called_once() # Ensure the failing method was called
# --- _find_similar_sessions_impl Tests ---
@pytest.mark.asyncio # Mark as async
async def test_find_similar_sessions_success(self, mock_chroma_client_thinking):
"""Test finding similar sessions successfully."""
mock_client, mock_thoughts_collection, mock_sessions_collection = mock_chroma_client_thinking
query = "Project planning"
threshold = 0.6
n_results = 5 # Test with default n_results
# --- Setup Mocks ---
# Mock the .get() call on the THOUGHTS collection to return session IDs
mock_thoughts_collection.get.return_value = {
"ids": ["t1", "t2", "t3"],
"metadatas": [
{"session_id": "session_abc"},
{"session_id": "session_xyz"},
{"session_id": "session_abc"}, # Duplicate ID is fine for get
],
# Other fields like documents are not needed for this specific mock
}
# Mock the SESSIONS collection query response
mock_sessions_collection.query.return_value = {
"ids": [["session_abc", "session_xyz"]],
"distances": [[0.3, 0.7]], # Similarities: 0.7, 0.3 (session_xyz filtered)
"metadatas": [[{"summary": "Plan A"}, {"summary": "Plan B"}]],
"documents": [[None], [None]],
}
# Mock the SESSIONS collection .get() call used for embedding new sessions
# Assume session_abc and session_xyz already exist
mock_sessions_collection.get.return_value = {"ids": ["session_abc", "session_xyz"]}
# Mock the internal call to _get_session_summary_impl for the final result assembly
with patch("src.chroma_mcp.tools.thinking_tools._get_session_summary_impl") as mock_get_summary:
# Configure mock_get_summary to return a successful result for session_abc
mock_summary_result_list = [
mcp_types.TextContent(
type="text",
text=json.dumps(
{
"session_id": "session_abc",
"session_thoughts": [{"id": "t1", "content": "plan A step 1"}],
"total_thoughts_in_session": 1,
}
),
)
]
mock_get_summary.return_value = mock_summary_result_list
# --- Act ---
input_model = FindSimilarSessionsInput(query=query, threshold=threshold)
result = await _find_similar_sessions_impl(input_model)
# --- Assert --- using MODIFIED helper
result_data = assert_successful_json_list_result(result)
# Assertions on mocks - Expect THOUGHTS first, then SESSIONS
mock_client.get_collection.assert_any_call(name=THOUGHTS_COLLECTION, embedding_function=ANY)
mock_client.get_or_create_collection.assert_any_call(name=SESSIONS_COLLECTION, embedding_function=ANY)
mock_sessions_collection.query.assert_called_once() # Check if query was made
# Assert query on the sessions collection mock
mock_sessions_collection.query.assert_called_once_with(
query_texts=[query], n_results=n_results, include=["metadatas", "distances"]
)
# Assertions on result data
assert len(result_data.get("similar_sessions", [])) == 1 # session_xyz filtered
assert result_data["similar_sessions"][0]["session_id"] == "session_abc"
assert result_data["similar_sessions"][0]["similarity_score"] == pytest.approx(0.7)
# Check that the summary data from the mocked _get_session_summary_impl is present
assert len(result_data["similar_sessions"][0].get("session_thoughts", [])) == 1
assert result_data["similar_sessions"][0]["session_thoughts"][0]["content"] == "plan A step 1"
@pytest.mark.asyncio # Mark as async
async def test_find_similar_sessions_collection_not_found(self, mock_chroma_client_thinking):
"""Test finding similar sessions when the sessions collection doesn't exist."""
mock_client, _, mock_sessions_collection = mock_chroma_client_thinking
# Configure get_collection to raise error specifically for SESSIONS_COLLECTION
def get_collection_side_effect(name, embedding_function=None):
if name == THOUGHTS_COLLECTION:
return MagicMock() # Return a dummy thoughts collection
elif name == SESSIONS_COLLECTION:
raise ValueError(f"Collection {SESSIONS_COLLECTION} does not exist.") # Simulate Chroma error
else:
raise ValueError(f"Unexpected collection {name}")
mock_client.get_collection.side_effect = get_collection_side_effect
# ACT
input_model = FindSimilarSessionsInput(query="any query")
result = await _find_similar_sessions_impl(input_model)
# ASSERT: Expect success list with empty results if THOUGHTS collection missing
result_data = assert_successful_json_list_result(result)
assert result_data.get("similar_sessions") == []
assert result_data.get("total_found") == 0
@pytest.mark.asyncio
async def test_find_similar_sessions_sessions_collection_not_found(self, mock_chroma_client_thinking):
"""Test finding similar sessions when the required SESSIONS collection doesn't exist."""
mock_client, mock_thoughts_collection, _ = mock_chroma_client_thinking
# Mock thoughts collection get to return some session IDs
mock_thoughts_collection.get.return_value = {"metadatas": [{"session_id": "s1"}]}
# Configure side effects for get_collection and get_or_create_collection
def get_collection_specific_side_effect(name, embedding_function=None):
if name == THOUGHTS_COLLECTION:
return mock_thoughts_collection
# For SESSIONS_COLLECTION, if get_collection is called directly and it doesn't exist
raise ValueError(f"Collection {name} does not exist. From get_collection direct call.")
def get_or_create_specific_side_effect(name, embedding_function=None):
if name == THOUGHTS_COLLECTION:
return mock_thoughts_collection # Thoughts collection is assumed to be fine
elif name == SESSIONS_COLLECTION:
# This is the crucial part for the test. Simulate that get_or_create_collection
# fails because the underlying get (or create) indicates non-existence in a way
# that the impl function expects for this specific McpError.
raise ValueError(f"Collection {SESSIONS_COLLECTION} does not exist.") # CORRECTED ERROR STRING
return MagicMock()
mock_client.get_collection.side_effect = get_collection_specific_side_effect
mock_client.get_or_create_collection.side_effect = get_or_create_specific_side_effect
input_model = FindSimilarSessionsInput(query="any query")
# Expect the specific McpError raised by _find_similar_sessions_impl
with assert_raises_mcp_error(f"Tool Error: Collection '{SESSIONS_COLLECTION}' not found"):
await _find_similar_sessions_impl(input_model)
# Verify calls that should have happened before the error
mock_client.get_collection.assert_any_call(name=THOUGHTS_COLLECTION, embedding_function=ANY)
mock_client.get_or_create_collection.assert_any_call(name=SESSIONS_COLLECTION, embedding_function=ANY)
# +++ Start New Coverage Tests +++
@pytest.mark.asyncio
async def test_sequential_thinking_get_create_collection_error(self, mock_chroma_client_thinking):
"""Test error during get_or_create_collection in sequential thinking."""
mock_client, _, _ = mock_chroma_client_thinking
error_message = "DB connection failed"
mock_client.get_or_create_collection.side_effect = Exception(error_message)
input_model = SequentialThinkingInput(thought="t", thought_number=1, total_thoughts=1)
# Update expected error message to match the wrapped one
expected_error = f"Could not access thinking collection: {error_message}"
with assert_raises_mcp_error(expected_error):
await _sequential_thinking_impl(input_model)
# Assert the failing mock was called
mock_client.get_or_create_collection.assert_called_once_with(name=THOUGHTS_COLLECTION, embedding_function=ANY)
@pytest.mark.asyncio
async def test_sequential_thinking_collection_add_error(self, mock_chroma_client_thinking):
"""Test error during collection.add in sequential thinking."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
error_message = "Failed to add document"
mock_collection.add.side_effect = ValueError(error_message)
input_model = SequentialThinkingInput(thought="t", thought_number=1, total_thoughts=1)
with assert_raises_mcp_error(f"ChromaDB Error adding thought: {error_message}"):
await _sequential_thinking_impl(input_model)
@pytest.mark.asyncio
async def test_sequential_thinking_get_prev_thoughts_error(self, mock_chroma_client_thinking):
"""Test non-critical error during collection.get for previous thoughts."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
mock_collection.get.side_effect = Exception("Failed to get previous")
# Call with thought_number > 1 to trigger the .get call
input_model = SequentialThinkingInput(thought="t2", thought_number=2, total_thoughts=2, session_id="s1")
# ACT: Should still succeed despite the get error, just log a warning
result = await _sequential_thinking_impl(input_model)
result_data = assert_successful_json_list_result(result)
# ASSERT: Check that the thought was added, but count is 0
mock_collection.add.assert_called_once()
mock_collection.get.assert_called_once() # Ensure get was called
assert result_data["previous_thoughts_count"] == 0
assert result_data["session_id"] == "s1"
@pytest.mark.asyncio
async def test_find_similar_thoughts_query_error(self, mock_chroma_client_thinking):
"""Test error during collection.query in find similar thoughts."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
error_message = "Invalid query syntax"
mock_collection.query.side_effect = ValueError(error_message)
input_model = FindSimilarThoughtsInput(query="q")
with assert_raises_mcp_error(f"ChromaDB Query Error: {error_message}"):
await _find_similar_thoughts_impl(input_model)
@pytest.mark.asyncio
async def test_find_similar_thoughts_generic_error(self, mock_chroma_client_thinking):
"""Test generic error during find similar thoughts."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
error_message = "Something unexpected happened"
mock_collection.query.side_effect = Exception(error_message)
input_model = FindSimilarThoughtsInput(query="q")
with assert_raises_mcp_error(
f"Tool Error: An unexpected error occurred while finding similar thoughts. Details: {error_message}"
):
await _find_similar_thoughts_impl(input_model)
@pytest.mark.asyncio
async def test_get_session_summary_get_error(self, mock_chroma_client_thinking):
"""Test error during collection.get in get session summary."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
error_message = "Invalid filter for get"
mock_collection.get.side_effect = ValueError(error_message)
input_model = GetSessionSummaryInput(session_id="s1")
with assert_raises_mcp_error(f"ChromaDB Get Error: {error_message}"):
await _get_session_summary_impl(input_model)
@pytest.mark.asyncio
async def test_find_similar_sessions_thoughts_get_error(self, mock_chroma_client_thinking):
"""Test error getting thoughts collection in find similar sessions."""
mock_client, mock_thoughts_collection, _ = mock_chroma_client_thinking
error_message = "Cannot access thoughts DB"
mock_thoughts_collection.get.side_effect = Exception(error_message)
input_model = FindSimilarSessionsInput(query="q")
with assert_raises_mcp_error(f"ChromaDB Error accessing thoughts collection: {error_message}"):
await _find_similar_sessions_impl(input_model)
@pytest.mark.asyncio
async def test_find_similar_sessions_sessions_get_error(self, mock_chroma_client_thinking):
"""Test error getting sessions collection in find similar sessions."""
mock_client, mock_thoughts_collection, mock_sessions_collection = mock_chroma_client_thinking
error_message = "Cannot access sessions DB"
# Mock thoughts collection successfully
mock_thoughts_collection.get.return_value = {"metadatas": [{"session_id": "s1"}]}
# Mock sessions collection get error
mock_sessions_collection.get.side_effect = Exception(error_message)
input_model = FindSimilarSessionsInput(query="q")
# This error happens during the check for existing sessions
with assert_raises_mcp_error(f"ChromaDB Error updating sessions: {error_message}"):
await _find_similar_sessions_impl(input_model)
@pytest.mark.asyncio
async def test_find_similar_sessions_embed_error(self, mock_chroma_client_thinking):
"""Test error embedding/adding sessions in find similar sessions."""
mock_client, mock_thoughts_collection, mock_sessions_collection = mock_chroma_client_thinking
error_message = "Embedding failed"
# Mock thoughts collection successfully
mock_thoughts_collection.get.return_value = {"metadatas": [{"session_id": "s1"}]}
# Mock sessions collection get returns empty (to trigger add)
mock_sessions_collection.get.return_value = {"ids": []}
# Mock sessions collection add error
mock_sessions_collection.add.side_effect = Exception(error_message)
# Mock the internal call to _get_session_summary_impl for embedding
with patch("src.chroma_mcp.tools.thinking_tools._get_session_summary_impl") as mock_get_summary:
# Return a list containing a TextContent object, not CallToolResult
mock_summary_result_list = [
mcp_types.TextContent(type="text", text=json.dumps({"session_thoughts": [{"content": "Summary"}]}))
]
mock_get_summary.return_value = mock_summary_result_list
input_model = FindSimilarSessionsInput(query="q")
with assert_raises_mcp_error(f"ChromaDB Error updating sessions: {error_message}"):
await _find_similar_sessions_impl(input_model)
@pytest.mark.asyncio
async def test_find_similar_sessions_query_error(self, mock_chroma_client_thinking):
"""Test error during query on sessions collection in find similar sessions."""
mock_client, mock_thoughts_collection, mock_sessions_collection = mock_chroma_client_thinking
error_message = "Session query failed"
# Mock thoughts collection successfully
mock_thoughts_collection.get.return_value = {"metadatas": [{"session_id": "s1"}]}
# Mock sessions collection get returns existing
mock_sessions_collection.get.return_value = {"ids": ["s1"]}
# Mock sessions collection query error
mock_sessions_collection.query.side_effect = ValueError(error_message)
input_model = FindSimilarSessionsInput(query="q")
with assert_raises_mcp_error(f"ChromaDB Query Error on sessions: {error_message}"):
await _find_similar_sessions_impl(input_model)
@pytest.mark.asyncio
async def test_find_similar_sessions_invalid_threshold(self, mock_chroma_client_thinking):
"""Test finding similar sessions with an invalid threshold value."""
mock_client, mock_collection, _ = mock_chroma_client_thinking
input_data_low = FindSimilarSessionsInput(query="test", threshold=-0.1)
input_data_high = FindSimilarSessionsInput(query="test", threshold=1.1)
expected_msg = "Threshold must be between 0.0 and 1.0"
with assert_raises_mcp_error(expected_msg):
await _find_similar_sessions_impl(input_data_low)
with assert_raises_mcp_error(expected_msg):
await _find_similar_sessions_impl(input_data_high)
# +++ End New Coverage Tests +++