Chroma MCP Server

"""Tests for collection management tools.""" import pytest import uuid import json from typing import Dict, Any, List, Optional from unittest.mock import patch, MagicMock, AsyncMock, ANY, call from contextlib import contextmanager from mcp import types from mcp.shared.exceptions import McpError from mcp.types import INVALID_PARAMS, INTERNAL_ERROR, ErrorData # Import specific errors if needed, or rely on ValidationError/Exception from src.chroma_mcp.utils.errors import ValidationError # --- Add direct import for setting global config --- from src.chroma_mcp import utils as chroma_utils # --- Add import for reset_client --- from src.chroma_mcp.utils import chroma_client as client_utils from src.chroma_mcp.tools.collection_tools import ( _reconstruct_metadata, # Keep helper if used _create_collection_impl, _create_collection_with_metadata_impl, _list_collections_impl, _get_collection_impl, _rename_collection_impl, _delete_collection_impl, _peek_collection_impl, ) # Import Pydantic models used by the tools from src.chroma_mcp.tools.collection_tools import ( CreateCollectionInput, CreateCollectionWithMetadataInput, ListCollectionsInput, GetCollectionInput, RenameCollectionInput, DeleteCollectionInput, PeekCollectionInput, ) # Correct import for get_collection_settings from src.chroma_mcp.utils.config import get_collection_settings from src.chroma_mcp.types import ChromaClientConfig # Add import for direct patching from src.chroma_mcp.tools import collection_tools DEFAULT_SIMILARITY_THRESHOLD = 0.7 # Mock embedding function classes used in chroma_client class MockEmbeddingFunction: def __init__(self, name="default"): # Add name for identification self.name = name def __call__(self, input: List[str]): # Simple mock implementation return [[0.1] * 10 for _ in input] # Use 'input' here too # --- Fixtures --- @pytest.fixture(autouse=True) def mock_logger(): """Auto-mock the logger.""" with patch("src.chroma_mcp.tools.collection_tools.get_logger") as mock_get_logger: mock_log_instance = MagicMock() mock_get_logger.return_value = mock_log_instance yield mock_log_instance @pytest.fixture def mock_chroma_client(): """Mock the ChromaDB client instance.""" mock_client = MagicMock() # Mock the create_collection method specifically mock_client.create_collection = MagicMock() # Mock get_collection, list_collections etc. as needed for other tests mock_client.get_collection = MagicMock() mock_client.list_collections = MagicMock(return_value=[]) mock_client.delete_collection = MagicMock() with patch("chroma_mcp.tools.collection_tools.get_chroma_client", return_value=mock_client): yield mock_client @pytest.fixture def mock_embedding_functions(): """Mock the get_embedding_function lookup.""" mock_functions = { "default": MockEmbeddingFunction(name="default"), "fast": MockEmbeddingFunction(name="fast"), "accurate": MockEmbeddingFunction(name="accurate"), "openai": MockEmbeddingFunction(name="openai"), } def mock_lookup(name): func = mock_functions.get(name.lower()) if func is None: raise McpError(ErrorData(code=INVALID_PARAMS, message=f"Unknown embedding function: {name}")) return func # Patch where get_embedding_function is looked up in the tools module with patch("src.chroma_mcp.tools.collection_tools.get_embedding_function", side_effect=mock_lookup) as mock_getter: yield mock_getter # Yield the mock getter itself for assertions @pytest.fixture def mock_server_config(): """Mock the server configuration retrieval.""" # Make this instance unique per test potentially mock_config = ChromaClientConfig(client_type="ephemeral", embedding_function_name="default") original_config = chroma_utils._global_client_config chroma_utils._global_client_config = mock_config try: yield mock_config # Yield the config object to modify it in tests finally: # Ensure global config is reset after test chroma_utils._global_client_config = original_config @pytest.fixture def mock_collection_settings(): """Mock retrieval of default collection settings.""" default_settings = {"hnsw:space": "cosine"} with patch("chroma_mcp.tools.collection_tools.get_collection_settings", return_value=default_settings): yield default_settings @pytest.fixture def mock_validate_name(): """Mock collection name validation.""" with patch("chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate: yield mock_validate # --- Helper Functions --- def assert_successful_json_result( result: List[types.TextContent], expected_data: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Asserts the tool result is a successful list containing valid JSON, returning the parsed data.""" assert isinstance(result, list) assert len(result) > 0, "Result list cannot be empty for successful JSON result." content_item = result[0] assert isinstance(content_item, types.TextContent), f"Expected TextContent, got {type(content_item)}" assert content_item.type == "text", f"Expected content type 'text', got '{content_item.type}'" assert content_item.text is not None, "Text content cannot be None for JSON result." try: parsed_data = json.loads(content_item.text) assert isinstance(parsed_data, dict), f"Parsed JSON is not a dictionary, got {type(parsed_data)}" except (json.JSONDecodeError, AssertionError) as e: pytest.fail(f"Result content is not valid JSON: {e}\nContent: {content_item.text}") if expected_data is not None: assert parsed_data == expected_data, f"Parsed JSON data mismatch. Expected: {expected_data}, Got: {parsed_data}" return parsed_data # Rework assert_error_result as a context manager for future use @contextmanager def assert_raises_mcp_error(expected_message: str): """Asserts that McpError is raised and optionally checks the error message.""" 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 # Apply usefixtures to the whole class again # Apply the new reset fixture as well # Restore mock_server_config to class level @pytest.mark.usefixtures("mock_embedding_functions", "mock_server_config") class TestCollectionTools: """Test cases for collection management tools.""" # --- _create_collection_impl Tests --- @pytest.mark.asyncio # REMOVE fixtures, use local patching async def test_create_collection_uses_default_ef_from_config(self): """Test that _create_collection_impl uses the EF name from server config.""" collection_name = "test-default-ef" # --- Mocks --- # mock_collection = MagicMock() mock_collection.name = collection_name mock_collection.id = "uuid-123" default_settings = get_collection_settings() # Get defaults stored_metadata = {f"chroma:setting:{k.replace(':', '_')}": v for k, v in default_settings.items()} mock_collection.metadata = stored_metadata mock_collection.count.return_value = 0 mock_client_instance = MagicMock() mock_client_instance.create_collection.return_value = mock_collection mock_ef_instance = MockEmbeddingFunction(name="default") # --- Patching --- # with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=mock_client_instance ) as mock_get_client, patch( "src.chroma_mcp.tools.collection_tools.get_embedding_function", return_value=mock_ef_instance ) as mock_get_ef, patch( "src.chroma_mcp.tools.collection_tools.get_collection_settings", return_value=default_settings ) as mock_get_settings, ): # --- Act --- # input_data = CreateCollectionInput(collection_name=collection_name) result = await _create_collection_impl(input_data) # --- Assert --- # mock_validate_name.assert_called_once_with(collection_name) mock_get_client.assert_called_once() mock_get_ef.assert_called_once_with("default") # Assumes default config mock_get_settings.assert_called_once() # Ensure settings were fetched expected_metadata_passed = {f"chroma:setting:{k.replace(':', '_')}": v for k, v in default_settings.items()} mock_client_instance.create_collection.assert_called_once_with( name=collection_name, metadata=expected_metadata_passed, embedding_function=mock_ef_instance, get_or_create=False, ) assert len(result) == 1 assert result[0].type == "text" # Add assertions on result content if necessary @pytest.mark.asyncio # REMOVE fixtures, use local patching async def test_create_collection_falls_back_to_default_ef(self): """Test _create_collection_impl uses 'default' EF name (implicitly).""" collection_name = "test-fallback-ef" # --- Mocks --- # mock_collection = MagicMock() mock_collection.name = collection_name mock_collection.id = "uuid-123" default_settings = get_collection_settings() # Get defaults stored_metadata = {f"chroma:setting:{k.replace(':', '_')}": v for k, v in default_settings.items()} mock_collection.metadata = stored_metadata mock_collection.count.return_value = 0 mock_client_instance = MagicMock() mock_client_instance.create_collection.return_value = mock_collection mock_ef_instance = MockEmbeddingFunction(name="default") # --- Patching --- # with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=mock_client_instance ) as mock_get_client, patch( "src.chroma_mcp.tools.collection_tools.get_embedding_function", return_value=mock_ef_instance ) as mock_get_ef, patch( "src.chroma_mcp.tools.collection_tools.get_collection_settings", return_value=default_settings ) as mock_get_settings, ): # --- Act --- # input_data = CreateCollectionInput(collection_name=collection_name) await _create_collection_impl(input_data) # --- Assert --- # mock_validate_name.assert_called_once_with(collection_name) mock_get_client.assert_called_once() mock_get_ef.assert_called_once_with("default") # Main assertion: default was used mock_get_settings.assert_called_once() expected_metadata_passed = {f"chroma:setting:{k.replace(':', '_')}": v for k, v in default_settings.items()} mock_client_instance.create_collection.assert_called_once_with( name=collection_name, metadata=expected_metadata_passed, embedding_function=mock_ef_instance, get_or_create=False, ) @pytest.mark.asyncio # REMOVE fixtures, use local patching async def test_create_collection_success(self): """Test successful collection creation.""" collection_name = "test_create_new" mock_collection_id = str(uuid.uuid4()) # --- Mocks --- # mock_collection = MagicMock() mock_collection.name = collection_name mock_collection.id = mock_collection_id default_settings = get_collection_settings() # Get defaults metadata_stored_by_chroma = {f"chroma:setting:{k.replace(':', '_')}": v for k, v in default_settings.items()} mock_collection.metadata = metadata_stored_by_chroma mock_collection.count.return_value = 0 mock_client_instance = MagicMock() mock_client_instance.create_collection.return_value = mock_collection mock_ef_instance = MockEmbeddingFunction(name="default") # --- Patching --- # with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=mock_client_instance ) as mock_get_client, patch( "src.chroma_mcp.tools.collection_tools.get_embedding_function", return_value=mock_ef_instance ) as mock_get_ef, patch( "src.chroma_mcp.tools.collection_tools.get_collection_settings", return_value=default_settings ) as mock_get_settings, ): # --- Act --- # input_model = CreateCollectionInput(collection_name=collection_name) result_list = await _create_collection_impl(input_model) # --- Assert --- # mock_validate_name.assert_called_once_with(collection_name) mock_get_client.assert_called_once() mock_get_ef.assert_called_once_with("default") mock_get_settings.assert_called_once() mock_client_instance.create_collection.assert_called_once() call_args, call_kwargs = mock_client_instance.create_collection.call_args assert call_kwargs["name"] == collection_name assert call_kwargs["metadata"] == metadata_stored_by_chroma assert call_kwargs["get_or_create"] is False assert call_kwargs["embedding_function"] is mock_ef_instance result_data = assert_successful_json_result(result_list) assert result_data.get("name") == collection_name assert result_data.get("id") == mock_collection_id assert "metadata" in result_data assert "settings" in result_data["metadata"] expected_settings_with_colons = {k.replace("_", ":"): v for k, v in default_settings.items()} assert result_data["metadata"]["settings"] == expected_settings_with_colons assert result_data.get("count") == 0 @pytest.mark.asyncio # REMOVE mock_chroma_client fixture # No fixtures needed beyond class level async def test_create_collection_invalid_name(self): """Test collection name validation failure within the implementation.""" invalid_name = "invalid-" error_msg = "Invalid collection name" # Patch validator and client method (to check not called) with ( patch( "src.chroma_mcp.tools.collection_tools.validate_collection_name", side_effect=ValidationError(error_msg) ) as mock_validate_name, patch("chromadb.api.client.Client.create_collection") as mock_create_collection, ): # --- Act & Assert --- input_model = CreateCollectionInput(collection_name=invalid_name) with assert_raises_mcp_error(f"Validation Error: {error_msg}"): await _create_collection_impl(input_model) mock_validate_name.assert_called_once_with(invalid_name) mock_create_collection.assert_not_called() # --- _peek_collection_impl Tests --- @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_peek_collection_success(self): """Test successful peeking into a collection with a specific limit.""" collection_name = "test_peek_exists" limit = 3 expected_peek_result = { "ids": ["id1", "id2", "id3"], "documents": ["doc1", "doc2", "doc3"], "metadatas": [{"m": 1}, {"m": 2}, {"m": 3}], "embeddings": None, } # Local Mocks local_mock_collection = MagicMock() local_mock_collection.peek.return_value = expected_peek_result local_mock_client = MagicMock() local_mock_client.get_collection.return_value = local_mock_collection # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = PeekCollectionInput(collection_name=collection_name, limit=limit) result = await _peek_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) local_mock_collection.peek.assert_called_once_with(limit=limit) assert_successful_json_result(result, expected_peek_result) @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_peek_collection_success_default_limit(self): """Test successful peeking using the default limit (10).""" collection_name = "test_peek_default" expected_peek_result = {"ids": ["default_id"]} # Local Mocks local_mock_collection = MagicMock() local_mock_collection.peek.return_value = expected_peek_result local_mock_client = MagicMock() local_mock_client.get_collection.return_value = local_mock_collection # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = PeekCollectionInput(collection_name=collection_name) # Use default limit result = await _peek_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) local_mock_collection.peek.assert_called_once_with(limit=10) # Check default limit value assert_successful_json_result(result, expected_peek_result) @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_peek_collection_validation_error(self): """Test collection name validation failure for peek.""" invalid_name = "peek-invalid--" error_msg = "Bad name for peek" # Patch locally (only validator needed, client shouldn't be called) with ( patch( "src.chroma_mcp.tools.collection_tools.validate_collection_name", side_effect=ValidationError(error_msg) ) as mock_validate, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_client, ): input_model = PeekCollectionInput(collection_name=invalid_name) # Expect ValidationError message format with assert_raises_mcp_error(f"Validation Error: {error_msg}"): await _peek_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(invalid_name) mock_get_client.assert_not_called() # Client should not be fetched if validation fails @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_peek_collection_not_found_error(self): """Test peek when the collection is not found (error from get_collection).""" collection_name = "peek_not_found" error_msg = f"Collection {collection_name} does not exist." # Local mock client configured to raise error local_mock_client = MagicMock() local_mock_client.get_collection.side_effect = ValueError(error_msg) # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = PeekCollectionInput(collection_name=collection_name) # Expect the specific 'not found' message from the ValueError handler with assert_raises_mcp_error(f"Tool Error: Collection '{collection_name}' not found."): await _peek_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_peek_collection_get_other_value_error(self): """Test other ValueError from get_collection during peek.""" collection_name = "peek_get_val_err" error_msg = "Get value error during peek" # Local mock client configured to raise error local_mock_client = MagicMock() local_mock_client.get_collection.side_effect = ValueError(error_msg) # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = PeekCollectionInput(collection_name=collection_name) # Expect the specific 'Problem accessing collection' message from the ValueError handler with assert_raises_mcp_error( f"Tool Error: Problem accessing collection '{collection_name}'. Details: {error_msg}" ): await _peek_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_peek_collection_peek_exception(self): """Test handling of exception during the actual collection.peek call.""" collection_name = "peek_itself_fails" error_msg = "Peek failed internally" # Local Mocks local_mock_collection = MagicMock() local_mock_collection.peek.side_effect = Exception(error_msg) local_mock_client = MagicMock() local_mock_client.get_collection.return_value = local_mock_collection # Patch Locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = PeekCollectionInput(collection_name=collection_name) # Expect the generic exception handler's message format expected_error_msg = ( f"Tool Error: An unexpected error occurred peeking collection '{collection_name}'. Details: {error_msg}" ) with assert_raises_mcp_error(expected_error_msg): await _peek_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) # Check default limit for PeekCollectionInput is 10 local_mock_collection.peek.assert_called_once_with(limit=10) # --- _create_collection_with_metadata_impl Tests --- @pytest.mark.asyncio # Remove all fixtures except reset_chroma_client_cache (autouse=True) @pytest.mark.usefixtures() async def test_create_collection_with_metadata_success(self): # No fixture arguments needed """Test successful creation using the _with_metadata variant (as JSON string).""" collection_name = "test_create_with_meta_success" custom_metadata_dict = {"user_key": "user_value"} custom_metadata_json = json.dumps(custom_metadata_dict) mock_collection_id = str(uuid.uuid4()) # --- Setup Local Mocks --- # Mock for the collection object that create_collection will return created_collection_mock = MagicMock() created_collection_mock.name = collection_name created_collection_mock.id = mock_collection_id created_collection_mock.metadata = custom_metadata_dict created_collection_mock.count.return_value = 0 # Local mock for the Chroma client instance local_mock_client = MagicMock() local_mock_client.create_collection = MagicMock(return_value=created_collection_mock) # Define dummy EF class and instance class MockEF: def __call__(self, input): return [[0.0] * 10 for _ in input] local_mock_ef_instance = MockEF() # Local mock for the logger instance local_mock_logger_instance = MagicMock() input_model = CreateCollectionWithMetadataInput(collection_name=collection_name, metadata=custom_metadata_json) # --- Patch all dependencies locally --- with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as local_mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_embedding_function", return_value=local_mock_ef_instance ) as local_mock_get_ef, patch( "src.chroma_mcp.tools.collection_tools.get_logger", return_value=local_mock_logger_instance ) as mock_get_logger_call, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as local_mock_get_client, ): # Note: We still rely on get_server_config implicitly inside the code for the default EF name # --- Act --- (Call the implementation) result_list = await _create_collection_with_metadata_impl(input_model) # --- Assert Internal Calls --- local_mock_validate.assert_called_once_with(collection_name) local_mock_get_ef.assert_called_once_with("default") # Assuming default EF name mock_get_logger_call.assert_called_once() local_mock_get_client.assert_called_once() # Ensure client getter was called # Assert Logger Call (should pass) expected_log_message = f"Attempting to create collection '{collection_name}' with provided parsed metadata: {custom_metadata_dict}" local_mock_logger_instance.info.assert_any_call(expected_log_message) assert local_mock_logger_instance.error.call_count == 0 # --- Assert Actual Client Call on the LOCAL mock --- local_mock_client.create_collection.assert_called_once() # Check args passed to the local client mock call_args = local_mock_client.create_collection.call_args assert call_args.kwargs["name"] == collection_name assert call_args.kwargs["metadata"] == custom_metadata_dict assert call_args.kwargs["embedding_function"] is local_mock_ef_instance # --- Assert Result --- result_data = assert_successful_json_result(result_list) assert result_data.get("name") == collection_name assert result_data.get("id") == mock_collection_id @pytest.mark.asyncio # Patch create_collection via decorator, use local patch for validator @patch("chromadb.api.client.Client.create_collection") @pytest.mark.usefixtures("mock_chroma_client", "mock_server_config", "mock_embedding_functions") async def test_create_collection_with_metadata_validation_error( self, mock_create_collection, mock_chroma_client, mock_server_config, mock_embedding_functions ): """Test create with metadata validation error using local patch for validator.""" invalid_name = "meta-invalid--" error_msg = "Bad name from local patch" valid_metadata_json = '{"key": "value"}' # Patch the validator locally where it's imported in the implementation module with patch( "src.chroma_mcp.tools.collection_tools.validate_collection_name", side_effect=ValidationError(error_msg) ) as local_mock_validate: input_model = CreateCollectionWithMetadataInput(collection_name=invalid_name, metadata=valid_metadata_json) with assert_raises_mcp_error(f"Validation Error: {error_msg}"): await _create_collection_with_metadata_impl(input_model) local_mock_validate.assert_called_once_with(invalid_name) mock_create_collection.assert_not_called() # Client method should not be called @pytest.mark.asyncio async def test_create_collection_with_metadata_invalid_json(self): """Test create with metadata when the metadata string is invalid JSON.""" collection_name = "valid_json_test_coll" invalid_json_string = '{"key": "value", "unterminated' expected_error_substring = "Invalid JSON format for metadata field" # More general check # Patch only the client method to check it's not called with patch("chromadb.api.client.Client.create_collection") as mock_create_collection: input_model = CreateCollectionWithMetadataInput( collection_name=collection_name, metadata=invalid_json_string ) with assert_raises_mcp_error(expected_error_substring): await _create_collection_with_metadata_impl(input_model) mock_create_collection.assert_not_called() @pytest.mark.asyncio async def test_create_collection_with_metadata_already_exists(self): """Test create with metadata when collection already exists (ValueError from client).""" collection_name = "meta_exists_test" error_msg = f"Collection {collection_name} already exists." valid_metadata_json = '{"key": "value"}' # Patch validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_client, ): # Configure the mock client returned by get_chroma_client mock_client_instance = MagicMock() # Set the side effect directly on the mock client's method mock_client_instance.create_collection.side_effect = ValueError(error_msg) mock_get_client.return_value = mock_client_instance input_model = CreateCollectionWithMetadataInput( collection_name=collection_name, metadata=valid_metadata_json ) # Check for the specific error message wrapped by the implementation with assert_raises_mcp_error(f"Tool Error: Collection '{collection_name}' already exists."): await _create_collection_with_metadata_impl(input_model) mock_validate_name.assert_called_once_with(collection_name) # Assert the create_collection method was called on the mock client instance mock_client_instance.create_collection.assert_called_once() @pytest.mark.asyncio async def test_create_collection_with_metadata_unexpected_error(self): """Test create with metadata handling unexpected errors from client.""" collection_name = "meta_unexpected_err_test" error_msg = "Something else failed badly" valid_metadata_json = '{"key": "value"}' # Patch validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_client, ): # Configure the mock client returned by get_chroma_client mock_client_instance = MagicMock() # Set the side effect directly on the mock client's method mock_client_instance.create_collection.side_effect = Exception(error_msg) mock_get_client.return_value = mock_client_instance input_model = CreateCollectionWithMetadataInput( collection_name=collection_name, metadata=valid_metadata_json ) # Check for the generic error message wrapped by the implementation with assert_raises_mcp_error( f"Tool Error: An unexpected error occurred while creating collection '{collection_name}'. Details: {error_msg}" ): await _create_collection_with_metadata_impl(input_model) mock_validate_name.assert_called_once_with(collection_name) # Assert the create_collection method was called on the mock client instance mock_client_instance.create_collection.assert_called_once() # --- _list_collections_impl Tests --- @pytest.mark.asyncio # REMOVE Fixture - Use local patching async def test_list_collections_success_defaults(self): """Test successful default collection listing (no filters, no pagination).""" all_names = ["coll_a", "coll_b"] # Local Mock local_mock_client = MagicMock() local_mock_client.list_collections.return_value = all_names # Patch locally with patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client: input_model = ListCollectionsInput() result_list = await _list_collections_impl(input_model) # Assertions mock_get_client.assert_called_once() local_mock_client.list_collections.assert_called_once() local_mock_client.get_collection.assert_not_called() # Still shouldn't be called result_data = assert_successful_json_result(result_list) assert result_data.get("collection_names") == all_names assert result_data.get("total_count") == len(all_names) assert result_data.get("limit") == 0 assert result_data.get("offset") == 0 @pytest.mark.asyncio # REMOVE Fixture - Use local patching async def test_list_collections_with_filter_pagination(self): """Test listing with name filter and pagination.""" all_names = ["apple", "banana", "apricot", "avocado"] # Local Mock local_mock_client = MagicMock() local_mock_client.list_collections.return_value = all_names # Patch locally with patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client: input_model = ListCollectionsInput(limit=1, offset=1, name_contains="ap") result_list = await _list_collections_impl(input_model) # Assertions mock_get_client.assert_called_once() local_mock_client.list_collections.assert_called_once() local_mock_client.get_collection.assert_not_called() result_data = assert_successful_json_result(result_list) assert result_data.get("collection_names") == ["apricot"] assert result_data.get("total_count") == 2 assert result_data.get("limit") == 1 assert result_data.get("offset") == 1 @pytest.mark.skip( reason="got empty parameter set ['limit', 'offset', 'expected_error_msg'], function test_list_collections_validation_error at /Users/dominikus/git/nold-ai/chroma_mcp_server/tests/tools/test_collection_tools.py:745" ) @pytest.mark.asyncio @pytest.mark.parametrize( "limit, offset, expected_error_msg", [ # (-1, 0, "limit cannot be negative"), # Pydantic handles this # (0, -1, "offset cannot be negative"), # Pydantic handles this ], ids=[ # "negative_limit", # Removed as Pydantic validation # "negative_offset", # Removed as Pydantic validation ], ) # Remove fixture use here, test only Pydantic validation async def test_list_collections_validation_error( self, limit, offset, expected_error_msg # Remove mock_chroma_client_collections ): """Test input validation for limit and offset in list_collections.""" # This test correctly uses Pydantic's validation, but we need one for client errors # No specific validation error raised IN the function, Pydantic handles it. # Instead, we add a test for client-side errors. # REMOVED the original test body as Pydantic handles this. # The assert_raises_mcp_error needs to be outside the await for Pydantic validation with pytest.raises(Exception): # Pydantic raises its own ValidationError _ = ListCollectionsInput(limit=limit, offset=offset) # Keep the assertion structure for documentation, though it might not be hit # with assert_raises_mcp_error(expected_error_msg): # input_model = ListCollectionsInput(limit=limit, offset=offset) # _ = await _list_collections_impl(input_model) @pytest.mark.asyncio # Patch the actual client method to raise error async def test_list_collections_client_error(self): """Test handling of errors during client.list_collections.""" error_msg = "Client connection failed" # Patch list_collections on the *mock client* returned by get_chroma_client with patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_client: # Configure the mock client mock_client_instance = MagicMock() mock_client_instance.list_collections.side_effect = Exception(error_msg) mock_get_client.return_value = mock_client_instance input_model = ListCollectionsInput() with assert_raises_mcp_error(f"Tool Error: Error listing collections. Details: {error_msg}"): await _list_collections_impl(input_model) # Assert the method was called on the mock client mock_client_instance.list_collections.assert_called_once() # --- _get_collection_impl Tests --- @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_get_collection_success(self): """Test getting existing collection info.""" collection_name = "my_coll" mock_collection_id = "test-id-123" mock_metadata_stored = {"description": "test desc", "chroma:setting:hnsw_space": "l2"} mock_count = 42 # Ensure embeddings are None in the peek result if not expected mock_peek = {"ids": ["p1"], "documents": ["peek doc"], "embeddings": None} # Local Mocks local_mock_collection = MagicMock() local_mock_collection.name = collection_name local_mock_collection.id = mock_collection_id local_mock_collection.metadata = mock_metadata_stored local_mock_collection.count.return_value = mock_count local_mock_collection.peek.return_value = mock_peek local_mock_client = MagicMock() local_mock_client.get_collection.return_value = local_mock_collection # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = GetCollectionInput(collection_name=collection_name) result = await _get_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) local_mock_collection.count.assert_called_once() local_mock_collection.peek.assert_called_once_with(limit=5) result_data = assert_successful_json_result(result) assert result_data.get("name") == collection_name assert result_data.get("id") == mock_collection_id assert result_data.get("count") == mock_count assert result_data.get("metadata") == _reconstruct_metadata(mock_metadata_stored) assert result_data.get("sample_entries") == mock_peek @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_get_collection_not_found(self): """Test getting a non-existent collection (handled in impl).""" collection_name = "not_found_coll" error_message = f"Collection {collection_name} does not exist." # Local Mock Client local_mock_client = MagicMock() # Configure get_collection to raise the specific not found error local_mock_client.get_collection.side_effect = ValueError(error_message) # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = GetCollectionInput(collection_name=collection_name) # Expect the specific error message from the ValueError handler with assert_raises_mcp_error(f"Tool Error: Collection '{collection_name}' not found."): await _get_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_get_collection_unexpected_error(self): """Test handling of unexpected error during get collection.""" collection_name = "test_get_fail" error_msg = "Get failed unexpectedly" # Local Mock Client local_mock_client = MagicMock() # Configure get_collection to raise a generic Exception local_mock_client.get_collection.side_effect = Exception(error_msg) # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = GetCollectionInput(collection_name=collection_name) # Expect the generic error message from the final except block expected_error = f"Tool Error: An unexpected error occurred while getting collection '{collection_name}'. Details: {error_msg}" with assert_raises_mcp_error(expected_error): await _get_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_get_collection_validation_error(self): """Test handling of validation error for collection name.""" collection_name = "invalid--name" error_msg = "Invalid name format" # Patch locally (validator raises, client getter shouldn't be called) with ( patch( "src.chroma_mcp.tools.collection_tools.validate_collection_name", side_effect=ValidationError(error_msg) ) as mock_validate, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_client, ): input_model = GetCollectionInput(collection_name=collection_name) # Expect the specific Validation Error message with assert_raises_mcp_error(f"Validation Error: {error_msg}"): await _get_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_not_called() # get_chroma_client should not be called @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_get_collection_other_value_error(self): """Test handling of other ValueErrors from get_collection.""" collection_name = "test_get_value_err" error_msg = "Another value error during get" # Local Mock Client local_mock_client = MagicMock() # Configure get_collection to raise a different ValueError local_mock_client.get_collection.side_effect = ValueError(error_msg) # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = GetCollectionInput(collection_name=collection_name) # Expect the specific 'Invalid parameter' message from the ValueError handler expected_error = f"Tool Error: Invalid parameter getting collection. Details: {error_msg}" with assert_raises_mcp_error(expected_error): await _get_collection_impl(input_model) # Assertions mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) @pytest.mark.asyncio # REMOVE Fixtures - Use local patching async def test_get_collection_peek_error(self): """Test that get_collection returns info even if peek fails.""" collection_name = "test_peek_fail" mock_collection_id = "peek-fail-id" mock_metadata_stored = {"desc": "test"} mock_count = 5 peek_error_msg = "Peek failed" # Local Mocks local_mock_collection = MagicMock() local_mock_collection.name = collection_name local_mock_collection.id = mock_collection_id local_mock_collection.metadata = mock_metadata_stored local_mock_collection.count.return_value = mock_count # Configure peek to raise an error local_mock_collection.peek.side_effect = Exception(peek_error_msg) local_mock_client = MagicMock() # Configure get_collection to succeed and return the collection mock local_mock_client.get_collection.return_value = local_mock_collection # Patch locally with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=local_mock_client ) as mock_get_client, ): input_model = GetCollectionInput(collection_name=collection_name) result_list = await _get_collection_impl(input_model) # Assertions on internal calls mock_validate.assert_called_once_with(collection_name) mock_get_client.assert_called_once() local_mock_client.get_collection.assert_called_once_with(name=collection_name) local_mock_collection.count.assert_called_once() local_mock_collection.peek.assert_called_once_with(limit=5) # Assert success, but check the sample_entries field for the error result_data = assert_successful_json_result(result_list) assert result_data["name"] == collection_name assert result_data["count"] == mock_count assert "sample_entries" in result_data assert isinstance(result_data["sample_entries"], dict) assert "error" in result_data["sample_entries"] assert peek_error_msg in result_data["sample_entries"]["error"] assert result_data.get("id") == mock_collection_id # Add check for ID assert result_data.get("metadata") == _reconstruct_metadata(mock_metadata_stored) # Add check for metadata # --- _rename_collection_impl Tests --- @pytest.mark.asyncio # Remove fixture, patch manually inside # @pytest.mark.usefixtures("mock_validate_name") async def test_rename_collection_success(self): """Test successful collection renaming.""" original_name = "rename_me" new_name = "renamed_successfully" # Patch validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_client, ): # Configure mock client and the collection it returns mock_client_instance = MagicMock() mock_collection_instance = MagicMock() mock_collection_instance.modify.return_value = None # modify returns None on success mock_client_instance.get_collection.return_value = mock_collection_instance mock_get_client.return_value = mock_client_instance # --- Act --- input_model = RenameCollectionInput(collection_name=original_name, new_name=new_name) result = await _rename_collection_impl(input_model) # --- Assert --- mock_validate_name.assert_has_calls([call(original_name), call(new_name)]) mock_get_client.assert_called_once() mock_client_instance.get_collection.assert_called_once_with(name=original_name) mock_collection_instance.modify.assert_called_once_with(name=new_name) # Assert successful result message assert isinstance(result, list) assert len(result) == 1 assert isinstance(result[0], types.TextContent) assert f"Collection '{original_name}' successfully renamed to '{new_name}'." in result[0].text @pytest.mark.asyncio # Remove fixture, patch manually inside # @pytest.mark.usefixtures("mock_validate_name") async def test_rename_collection_original_not_found(self): """Test renaming when the original collection does not exist.""" original_name = "original_not_found" new_name = "new_name_irrelevant" # Patch validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_client, ): # Configure the mock client's get_collection method to raise the error mock_client_instance = MagicMock() mock_client_instance.get_collection.side_effect = ValueError(f"Collection {original_name} does not exist.") mock_get_client.return_value = mock_client_instance # --- Act --- input_model = RenameCollectionInput(collection_name=original_name, new_name=new_name) with assert_raises_mcp_error(f"Tool Error: Collection '{original_name}' not found."): await _rename_collection_impl(input_model) # --- Assert --- mock_validate_name.assert_has_calls([call(original_name), call(new_name)]) # Assert the get_collection method was called on the mock client mock_client_instance.get_collection.assert_called_once_with(name=original_name) @pytest.mark.asyncio # Remove fixture, patch manually inside # @pytest.mark.usefixtures("mock_validate_name") async def test_rename_collection_new_name_exists(self): """Test renaming when the new name already exists.""" original_name = "original_exists" new_name = "new_name_exists" # Patch validator, client.get_collection, and collection.modify with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch("chromadb.api.client.Client.get_collection") as mock_get_collection_method, patch( "chromadb.api.models.Collection.Collection.modify", # Patch the real modify side_effect=ValueError(f"Collection {new_name} already exists."), ) as mock_modify_method, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=MagicMock() ) as mock_get_client, ): # Add this patch # Configure the mock collection returned by get_collection mock_collection_instance = MagicMock() @pytest.mark.asyncio # Remove fixture, patch manually inside # @pytest.mark.usefixtures("mock_validate_name") async def test_rename_collection_unexpected_error(self): """Test unexpected error during rename (e.g., during modify).""" original_name = "rename_fail_orig" new_name = "rename_fail_new" error_msg = "Rename blew up" # Patch validator, client.get_collection, and collection.modify with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch("chromadb.api.client.Client.get_collection") as mock_get_collection_method, patch( "chromadb.api.models.Collection.Collection.modify", # Patch the real modify side_effect=Exception(error_msg), ) as mock_modify_method, patch( "src.chroma_mcp.tools.collection_tools.get_chroma_client", return_value=MagicMock() ) as mock_get_client, ): # Add this patch # Configure the mock collection returned by get_collection mock_collection_instance = MagicMock() @pytest.mark.asyncio # Remove fixture, patch manually inside # @pytest.mark.usefixtures("mock_validate_name") async def test_rename_collection_other_value_error(self): """Test handling of other ValueErrors during rename (e.g., from modify).""" original_name = "rename_val_err_orig" new_name = "rename_val_err_new" error_msg = "Some other value issue" # Patch validator AND get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate_name, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_chroma_client, ): # Configure mock client and the collection it returns mock_client_instance = MagicMock() mock_collection_instance = MagicMock() mock_collection_instance.modify.side_effect = ValueError(error_msg) mock_client_instance.get_collection.return_value = mock_collection_instance # Make get_chroma_client return our mock client mock_get_chroma_client.return_value = mock_client_instance input_model = RenameCollectionInput(collection_name=original_name, new_name=new_name) # Expect the ValueError block (the one that doesn't check specific messages) with assert_raises_mcp_error(f"Invalid parameter during rename. Details: {error_msg}"): await _rename_collection_impl(input_model) # Assertions mock_validate_name.assert_any_call(original_name) mock_validate_name.assert_any_call(new_name) mock_get_chroma_client.assert_called_once() # Assert get_collection was called on the *mock* client mock_client_instance.get_collection.assert_called_once_with(name=original_name) # Assert modify was called on the *mock* collection mock_collection_instance.modify.assert_called_once_with(name=new_name) # --- _delete_collection_impl Tests --- @pytest.mark.asyncio # REMOVE explicit fixture if applied at class level # @pytest.mark.usefixtures("mock_server_config") async def test_delete_collection_success(self): # Keep self only """Test successful collection deletion.""" collection_name = "delete_me" # Use nested patches for validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_chroma_client, ): # Configure mock client and its delete method mock_client_instance = MagicMock() mock_client_instance.delete_collection.return_value = None # Successful delete mock_get_chroma_client.return_value = mock_client_instance # --- Act --- input_model = DeleteCollectionInput(collection_name=collection_name) result = await _delete_collection_impl(input_model) # --- Assert --- mock_validate.assert_called_once_with(collection_name) mock_get_chroma_client.assert_called_once() # Assert delete was called on the mock client instance mock_client_instance.delete_collection.assert_called_once_with(name=collection_name) # Check plain text result assert isinstance(result, list) assert len(result) == 1 assert isinstance(result[0], types.TextContent) assert result[0].type == "text" assert result[0].text == f"Collection '{collection_name}' deleted successfully." # assert_successful_json_result(result, {"message": f"Collection '{collection_name}' deleted successfully."}) @pytest.mark.asyncio # REMOVE explicit fixture if applied at class level # @pytest.mark.usefixtures("mock_server_config") async def test_delete_collection_not_found(self): # Keep self only """Test deleting a non-existent collection (error from delete_collection).""" collection_name = "not_found_delete" error_message = f"Collection {collection_name} does not exist." # Use nested patches for validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_chroma_client, ): # Configure mock client to raise ValueError on delete mock_client_instance = MagicMock() mock_client_instance.delete_collection.side_effect = ValueError(error_message) mock_get_chroma_client.return_value = mock_client_instance # --- Act & Assert --- input_model = DeleteCollectionInput(collection_name=collection_name) with assert_raises_mcp_error(f"Tool Error: Collection '{collection_name}' not found."): await _delete_collection_impl(input_model) # Assert mocks mock_validate.assert_called_once_with(collection_name) mock_get_chroma_client.assert_called_once() mock_client_instance.delete_collection.assert_called_once_with(name=collection_name) @pytest.mark.asyncio # REMOVE explicit fixture if applied at class level # @pytest.mark.usefixtures("mock_server_config") async def test_delete_collection_unexpected_error(self): # Keep self only """Test unexpected error during the delete_collection call.""" collection_name = "delete_fail" error_msg = "Delete failed unexpectedly" # Use nested patches for validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_chroma_client, ): # Configure mock client to raise generic Exception on delete mock_client_instance = MagicMock() mock_client_instance.delete_collection.side_effect = Exception(error_msg) mock_get_chroma_client.return_value = mock_client_instance input_model = DeleteCollectionInput(collection_name=collection_name) # Expect the generic Exception handler message with assert_raises_mcp_error( f"An unexpected error occurred deleting collection '{collection_name}'. Details: {error_msg}" ): await _delete_collection_impl(input_model) # Assert mocks mock_validate.assert_called_once_with(collection_name) mock_get_chroma_client.assert_called_once() mock_client_instance.delete_collection.assert_called_once_with(name=collection_name) @pytest.mark.asyncio # REMOVE explicit fixture if applied at class level # @pytest.mark.usefixtures("mock_server_config") async def test_delete_collection_validation_error(self): """Test validation error (e.g., invalid name) before client call.""" collection_name = "invalid name!" error_msg = "Invalid collection name format" # Patch only the validator to raise error with ( patch( "src.chroma_mcp.tools.collection_tools.validate_collection_name", side_effect=ValidationError(error_msg) ) as mock_validate, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_chroma_client, ): input_model = DeleteCollectionInput(collection_name=collection_name) with assert_raises_mcp_error(f"Validation Error: {error_msg}"): await _delete_collection_impl(input_model) # Assert mocks mock_validate.assert_called_once_with(collection_name) mock_get_chroma_client.assert_not_called() # Client should not be fetched @pytest.mark.asyncio # REMOVE explicit fixture if applied at class level # @pytest.mark.usefixtures("mock_server_config") async def test_delete_collection_other_value_error(self): # Keep self only """Test handling of other ValueErrors during the delete_collection call.""" collection_name = "delete_val_err" error_msg = "Another delete value error" # Use nested patches for validator and get_chroma_client with ( patch("src.chroma_mcp.tools.collection_tools.validate_collection_name") as mock_validate, patch("src.chroma_mcp.tools.collection_tools.get_chroma_client") as mock_get_chroma_client, ): # Configure mock client to raise specific ValueError (not 'not found') mock_client_instance = MagicMock() mock_client_instance.delete_collection.side_effect = ValueError(error_msg) mock_get_chroma_client.return_value = mock_client_instance input_model = DeleteCollectionInput(collection_name=collection_name) # Expect the specific ValueError handler message (check the impl for exact wording) with assert_raises_mcp_error( f"Invalid parameter during delete operation for '{collection_name}'. Details: {error_msg}" ): await _delete_collection_impl(input_model) # Assert mocks mock_validate.assert_called_once_with(collection_name) mock_get_chroma_client.assert_called_once() mock_client_instance.delete_collection.assert_called_once_with(name=collection_name)