crewAI/tests/tools/couchbase_tool_test.py

import pytest
from unittest.mock import MagicMock, patch, ANY

# Mock the couchbase library before importing the tool
# This prevents ImportErrors if couchbase isn't installed in the test environment
mock_couchbase = MagicMock()
mock_couchbase.search = MagicMock()
mock_couchbase.cluster = MagicMock()
mock_couchbase.options = MagicMock()
mock_couchbase.vector_search = MagicMock()

# Simulate the structure needed for checks
mock_couchbase.cluster.Cluster = MagicMock()
mock_couchbase.options.SearchOptions = MagicMock()
mock_couchbase.vector_search.VectorQuery = MagicMock()
mock_couchbase.vector_search.VectorSearch = MagicMock()
mock_couchbase.search.SearchRequest = MagicMock() # Mock the class itself
mock_couchbase.search.SearchRequest.create = MagicMock() # Mock the class method

# Add necessary exception types if needed for testing error handling
class MockCouchbaseException(Exception):
    pass
mock_couchbase.exceptions = MagicMock()
mock_couchbase.exceptions.BucketNotFoundException = MockCouchbaseException
mock_couchbase.exceptions.ScopeNotFoundException = MockCouchbaseException
mock_couchbase.exceptions.CollectionNotFoundException = MockCouchbaseException
mock_couchbase.exceptions.IndexNotFoundException = MockCouchbaseException


import sys
sys.modules['couchbase'] = mock_couchbase
sys.modules['couchbase.search'] = mock_couchbase.search
sys.modules['couchbase.cluster'] = mock_couchbase.cluster
sys.modules['couchbase.options'] = mock_couchbase.options
sys.modules['couchbase.vector_search'] = mock_couchbase.vector_search
sys.modules['couchbase.exceptions'] = mock_couchbase.exceptions

# Now import the tool
from crewai_tools.tools.couchbase_tool.couchbase_tool import CouchbaseFTSVectorSearchTool

# --- Test Fixtures ---
@pytest.fixture(autouse=True)
def reset_global_mocks():
    """Reset call counts for globally defined mocks before each test."""
    # Reset the specific mock causing the issue
    mock_couchbase.vector_search.VectorQuery.reset_mock()
    # It's good practice to also reset other related global mocks
    # that might be called in your tests to prevent similar issues:
    mock_couchbase.vector_search.VectorSearch.from_vector_query.reset_mock()
    mock_couchbase.search.SearchRequest.create.reset_mock()

# Additional fixture to handle import pollution in full test suite
@pytest.fixture(autouse=True)
def ensure_couchbase_mocks():
    """Ensure that couchbase imports are properly mocked even when other tests have run first."""
    # This fixture ensures our mocks are in place regardless of import order
    original_modules = {}
    
    # Store any existing modules
    for module_name in ['couchbase', 'couchbase.search', 'couchbase.cluster', 'couchbase.options', 'couchbase.vector_search', 'couchbase.exceptions']:
        if module_name in sys.modules:
            original_modules[module_name] = sys.modules[module_name]
    
    # Ensure our mocks are active
    sys.modules['couchbase'] = mock_couchbase
    sys.modules['couchbase.search'] = mock_couchbase.search
    sys.modules['couchbase.cluster'] = mock_couchbase.cluster
    sys.modules['couchbase.options'] = mock_couchbase.options
    sys.modules['couchbase.vector_search'] = mock_couchbase.vector_search
    sys.modules['couchbase.exceptions'] = mock_couchbase.exceptions
    
    yield
    
    # Restore original modules if they existed
    for module_name, original_module in original_modules.items():
        if original_module is not None:
            sys.modules[module_name] = original_module

@pytest.fixture
def mock_cluster():
    cluster = MagicMock()
    bucket_manager = MagicMock()
    search_index_manager = MagicMock()
    bucket = MagicMock()
    scope = MagicMock()
    collection = MagicMock()
    scope_search_index_manager = MagicMock()

    # Setup mock return values for checks
    cluster.buckets.return_value = bucket_manager
    cluster.search_indexes.return_value = search_index_manager
    cluster.bucket.return_value = bucket
    bucket.scope.return_value = scope
    scope.collection.return_value = collection
    scope.search_indexes.return_value = scope_search_index_manager

    # Mock bucket existence check
    bucket_manager.get_bucket.return_value = True

    # Mock scope/collection existence check
    mock_scope_spec = MagicMock()
    mock_scope_spec.name = "test_scope"
    mock_collection_spec = MagicMock()
    mock_collection_spec.name = "test_collection"
    mock_scope_spec.collections = [mock_collection_spec]
    bucket.collections.return_value.get_all_scopes.return_value = [mock_scope_spec]

    # Mock index existence check
    mock_index_def = MagicMock()
    mock_index_def.name = "test_index"
    scope_search_index_manager.get_all_indexes.return_value = [mock_index_def]
    search_index_manager.get_all_indexes.return_value = [mock_index_def]

    return cluster

@pytest.fixture
def mock_embedding_function():
    # Simple mock embedding function
    # return lambda query: [0.1] * 10 # Example embedding vector
    return MagicMock(return_value=[0.1] * 10)

@pytest.fixture
def tool_config(mock_cluster, mock_embedding_function):
    return {
        "cluster": mock_cluster,
        "bucket_name": "test_bucket",
        "scope_name": "test_scope",
        "collection_name": "test_collection",
        "index_name": "test_index",
        "embedding_function": mock_embedding_function,
        "limit": 5,
        "embedding_key": "test_embedding",
        "scoped_index": True
    }

@pytest.fixture
def couchbase_tool(tool_config):
    # Patch COUCHBASE_AVAILABLE to True for these tests
    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
        tool = CouchbaseFTSVectorSearchTool(**tool_config)
        return tool

@pytest.fixture
def mock_search_iter():
    mock_iter = MagicMock()
    # Simulate search results with a 'fields' attribute
    mock_row1 = MagicMock()
    mock_row1.fields = {"id": "doc1", "text": "content 1", "test_embedding": [0.1]*10}
    mock_row2 = MagicMock()
    mock_row2.fields = {"id": "doc2", "text": "content 2", "test_embedding": [0.2]*10}
    mock_iter.rows.return_value = [mock_row1, mock_row2]
    return mock_iter

# --- Test Cases ---

def test_initialization_success(couchbase_tool, tool_config):
    """Test successful initialization with valid config."""
    assert couchbase_tool.cluster == tool_config["cluster"]
    assert couchbase_tool.bucket_name == "test_bucket"
    assert couchbase_tool.scope_name == "test_scope"
    assert couchbase_tool.collection_name == "test_collection"
    assert couchbase_tool.index_name == "test_index"
    assert couchbase_tool.embedding_function is not None
    assert couchbase_tool.limit == 5
    assert couchbase_tool.embedding_key == "test_embedding"
    assert couchbase_tool.scoped_index == True

    # Check if helper methods were called during init (via mocks in fixture)
    couchbase_tool.cluster.buckets().get_bucket.assert_called_once_with("test_bucket")
    couchbase_tool.cluster.bucket().collections().get_all_scopes.assert_called_once()
    couchbase_tool.cluster.bucket().scope().search_indexes().get_all_indexes.assert_called_once()

def test_initialization_missing_required_args(mock_cluster, mock_embedding_function):
    """Test initialization fails when required arguments are missing."""
    base_config = {
        "cluster": mock_cluster, "bucket_name": "b", "scope_name": "s",
        "collection_name": "c", "index_name": "i", "embedding_function": mock_embedding_function
    }
    required_keys = base_config.keys()

    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
        for key in required_keys:
            incomplete_config = base_config.copy()
            del incomplete_config[key]
            with pytest.raises(ValueError):
                 CouchbaseFTSVectorSearchTool(**incomplete_config)

def test_initialization_couchbase_unavailable():
    """Test behavior when couchbase library is not available."""
    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', False):
         with patch('click.confirm', return_value=False) as mock_confirm:
             with pytest.raises(ImportError, match="The 'couchbase' package is required"):
                 CouchbaseFTSVectorSearchTool(cluster=MagicMock(), bucket_name="b", scope_name="s",
                                        collection_name="c", index_name="i", embedding_function=MagicMock())
             mock_confirm.assert_called_once() # Ensure user was prompted

def test_run_success_scoped_index(couchbase_tool, mock_search_iter, tool_config, mock_embedding_function):
    """Test successful _run execution with a scoped index."""
    query = "find relevant documents"
    # expected_embedding = mock_embedding_function(query)

    # Mock the scope search method
    couchbase_tool._scope.search = MagicMock(return_value=mock_search_iter)
    # Mock the VectorQuery/VectorSearch/SearchRequest creation using runtime patching
    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.VectorQuery') as mock_vq, \
         patch('crewai_tools.tools.couchbase_tool.couchbase_tool.VectorSearch') as mock_vs, \
         patch('crewai_tools.tools.couchbase_tool.couchbase_tool.search.SearchRequest') as mock_sr, \
         patch('crewai_tools.tools.couchbase_tool.couchbase_tool.SearchOptions') as mock_so:
        
        # Set up the mock objects and their return values
        mock_vector_query = MagicMock()
        mock_vector_search = MagicMock()
        mock_search_req = MagicMock()
        mock_search_options = MagicMock()
        
        mock_vq.return_value = mock_vector_query
        mock_vs.from_vector_query.return_value = mock_vector_search
        mock_sr.create.return_value = mock_search_req
        mock_so.return_value = mock_search_options

        result = couchbase_tool._run(query=query)

        # Check embedding function call
        tool_config['embedding_function'].assert_called_once_with(query)

        # Check VectorQuery call
        mock_vq.assert_called_once_with(
            tool_config['embedding_key'], mock_embedding_function.return_value, tool_config['limit']
        )
        # Check VectorSearch call
        mock_vs.from_vector_query.assert_called_once_with(mock_vector_query)
        # Check SearchRequest creation
        mock_sr.create.assert_called_once_with(mock_vector_search)
        # Check SearchOptions creation
        mock_so.assert_called_once_with(limit=tool_config['limit'], fields=["*"])

        # Check that scope search was called correctly
        couchbase_tool._scope.search.assert_called_once_with(
            tool_config['index_name'],
            mock_search_req,
            mock_search_options
        )

    # Check cluster search was NOT called
    couchbase_tool.cluster.search.assert_not_called()

    # Check result format (simple check for JSON structure)
    assert '"id": "doc1"' in result
    assert '"id": "doc2"' in result
    assert result.startswith('[') # Should be valid JSON after concatenation

def test_run_success_global_index(tool_config, mock_search_iter, mock_embedding_function):
    """Test successful _run execution with a global (non-scoped) index."""
    tool_config['scoped_index'] = False
    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
        couchbase_tool = CouchbaseFTSVectorSearchTool(**tool_config)

    query = "find global documents"
    # expected_embedding = mock_embedding_function(query)

    # Mock the cluster search method
    couchbase_tool.cluster.search = MagicMock(return_value=mock_search_iter)
    # Mock the VectorQuery/VectorSearch/SearchRequest creation using runtime patching
    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.VectorQuery') as mock_vq, \
         patch('crewai_tools.tools.couchbase_tool.couchbase_tool.VectorSearch') as mock_vs, \
         patch('crewai_tools.tools.couchbase_tool.couchbase_tool.search.SearchRequest') as mock_sr, \
         patch('crewai_tools.tools.couchbase_tool.couchbase_tool.SearchOptions') as mock_so:
        
        # Set up the mock objects and their return values
        mock_vector_query = MagicMock()
        mock_vector_search = MagicMock()
        mock_search_req = MagicMock()
        mock_search_options = MagicMock()
        
        mock_vq.return_value = mock_vector_query
        mock_vs.from_vector_query.return_value = mock_vector_search
        mock_sr.create.return_value = mock_search_req
        mock_so.return_value = mock_search_options

        result = couchbase_tool._run(query=query)

        # Check embedding function call
        tool_config['embedding_function'].assert_called_once_with(query)

        # Check VectorQuery/Search call
        mock_vq.assert_called_once_with(
            tool_config['embedding_key'], mock_embedding_function.return_value, tool_config['limit']
        )
        mock_sr.create.assert_called_once_with(mock_vector_search)
        # Check SearchOptions creation
        mock_so.assert_called_once_with(limit=tool_config['limit'], fields=["*"])

        # Check that cluster search was called correctly
        couchbase_tool.cluster.search.assert_called_once_with(
            tool_config['index_name'],
            mock_search_req,
            mock_search_options
        )

    # Check scope search was NOT called
    couchbase_tool._scope.search.assert_not_called()

    # Check result format
    assert '"id": "doc1"' in result
    assert '"id": "doc2"' in result

def test_check_bucket_exists_fail(tool_config):
    """Test check for bucket non-existence."""
    mock_cluster = tool_config['cluster']
    mock_cluster.buckets().get_bucket.side_effect = mock_couchbase.exceptions.BucketNotFoundException("Bucket not found")

    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
      with pytest.raises(ValueError, match="Bucket test_bucket does not exist."):
          CouchbaseFTSVectorSearchTool(**tool_config)


def test_check_scope_exists_fail(tool_config):
    """Test check for scope non-existence."""
    mock_cluster = tool_config['cluster']
    # Simulate scope not being in the list returned
    mock_scope_spec = MagicMock()
    mock_scope_spec.name = "wrong_scope"
    mock_cluster.bucket().collections().get_all_scopes.return_value = [mock_scope_spec]

    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
        with pytest.raises(ValueError, match="Scope test_scope not found"):
            CouchbaseFTSVectorSearchTool(**tool_config)


def test_check_collection_exists_fail(tool_config):
    """Test check for collection non-existence."""
    mock_cluster = tool_config['cluster']
    # Simulate collection not being in the scope's list
    mock_scope_spec = MagicMock()
    mock_scope_spec.name = "test_scope"
    mock_collection_spec = MagicMock()
    mock_collection_spec.name = "wrong_collection"
    mock_scope_spec.collections = [mock_collection_spec] # Only has wrong collection
    mock_cluster.bucket().collections().get_all_scopes.return_value = [mock_scope_spec]

    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
        with pytest.raises(ValueError, match="Collection test_collection not found"):
             CouchbaseFTSVectorSearchTool(**tool_config)

def test_check_index_exists_fail_scoped(tool_config):
    """Test check for scoped index non-existence."""
    mock_cluster = tool_config['cluster']
    # Simulate index not being in the list returned by scope manager
    mock_cluster.bucket().scope().search_indexes().get_all_indexes.return_value = []

    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
        with pytest.raises(ValueError, match="Index test_index does not exist"):
            CouchbaseFTSVectorSearchTool(**tool_config)


def test_check_index_exists_fail_global(tool_config):
    """Test check for global index non-existence."""
    tool_config['scoped_index'] = False
    mock_cluster = tool_config['cluster']
    # Simulate index not being in the list returned by cluster manager
    mock_cluster.search_indexes().get_all_indexes.return_value = []

    with patch('crewai_tools.tools.couchbase_tool.couchbase_tool.COUCHBASE_AVAILABLE', True):
        with pytest.raises(ValueError, match="Index test_index does not exist"):
            CouchbaseFTSVectorSearchTool(**tool_config)