LangBot/tests/unit_tests/vector/test_vdb_base.py

"""Tests for VectorDatabase base class and SearchType enum."""

from __future__ import annotations

from unittest.mock import AsyncMock
import pytest

from langbot.pkg.vector.vdb import (
    SearchType,
    VectorDatabase,
    remember_bounded_mapping,
    remember_bounded_set,
)


class TestSearchType:
    """Tests for SearchType enum."""

    def test_search_type_values(self):
        """Test SearchType enum values."""
        assert SearchType.VECTOR.value == 'vector'
        assert SearchType.FULL_TEXT.value == 'full_text'
        assert SearchType.HYBRID.value == 'hybrid'

    def test_search_type_is_string_enum(self):
        """SearchType is a string enum."""
        assert isinstance(SearchType.VECTOR, str)
        assert SearchType.VECTOR == 'vector'

    def test_search_type_from_string(self):
        """Can create SearchType from string."""
        assert SearchType('vector') == SearchType.VECTOR
        assert SearchType('full_text') == SearchType.FULL_TEXT
        assert SearchType('hybrid') == SearchType.HYBRID


def test_runtime_cache_helpers_bound_mapping_and_set():
    mapping = {}
    values = set()

    for index in range(100):
        remember_bounded_mapping(mapping, str(index), object(), 8)
        remember_bounded_set(values, str(index), 8)

    assert len(mapping) == 8
    assert len(values) == 8


class TestVectorDatabaseAbstractMethods:
    """Tests for VectorDatabase abstract methods."""

    def test_vector_database_is_abstract(self):
        """VectorDatabase is abstract and cannot be instantiated directly."""
        with pytest.raises(TypeError):
            VectorDatabase()

    def test_abstract_methods_required(self):
        """Subclass must implement all abstract methods."""

        class IncompleteVectorDB(VectorDatabase):
            pass

        with pytest.raises(TypeError):
            IncompleteVectorDB()

    def test_supported_search_types_default(self):
        """Default supported_search_types returns [VECTOR]."""

        class MinimalVectorDB(VectorDatabase):
            async def add_embeddings(self, collection, ids, embeddings_list, metadatas, documents=None):
                pass

            async def search(
                self,
                collection,
                query_embedding,
                k=5,
                search_type='vector',
                query_text='',
                filter=None,
                vector_weight=None,
            ):
                pass

            async def delete_by_file_id(self, collection, file_id):
                pass

            async def delete_by_filter(self, collection, filter):
                pass

            async def get_or_create_collection(self, collection):
                pass

            async def delete_collection(self, collection):
                pass

        db = MinimalVectorDB()
        assert db.supported_search_types() == [SearchType.VECTOR]

    def test_list_by_filter_default_implementation(self):
        """list_by_filter has default implementation returning empty."""

        class MinimalVectorDB(VectorDatabase):
            async def add_embeddings(self, collection, ids, embeddings_list, metadatas, documents=None):
                pass

            async def search(
                self,
                collection,
                query_embedding,
                k=5,
                search_type='vector',
                query_text='',
                filter=None,
                vector_weight=None,
            ):
                pass

            async def delete_by_file_id(self, collection, file_id):
                pass

            async def delete_by_filter(self, collection, filter):
                pass

            async def get_or_create_collection(self, collection):
                pass

            async def delete_collection(self, collection):
                pass

        db = MinimalVectorDB()
        # list_by_filter should return empty list and -1 for total
        import asyncio

        result = asyncio.get_event_loop().run_until_complete(db.list_by_filter('test_collection'))
        assert result == ([], -1)


class TestVectorDatabaseInterface:
    """Tests for VectorDatabase interface contracts."""

    @pytest.fixture
    def mock_vector_db(self):
        """Create a minimal mock VectorDatabase for testing."""

        class MockVectorDB(VectorDatabase):
            def __init__(self):
                self.add_embeddings = AsyncMock()
                self.search = AsyncMock(
                    return_value={
                        'ids': [['id1', 'id2']],
                        'distances': [[0.1, 0.2]],
                        'metadatas': [[{'key': 'val1'}, {'key': 'val2'}]],
                    }
                )
                self.delete_by_file_id = AsyncMock()
                self.delete_by_filter = AsyncMock(return_value=5)
                self.get_or_create_collection = AsyncMock()
                self.delete_collection = AsyncMock()

            async def add_embeddings(self, collection, ids, embeddings_list, metadatas, documents=None):
                pass

            async def search(
                self,
                collection,
                query_embedding,
                k=5,
                search_type='vector',
                query_text='',
                filter=None,
                vector_weight=None,
            ):
                pass

            async def delete_by_file_id(self, collection, file_id):
                pass

            async def delete_by_filter(self, collection, filter):
                pass

            async def get_or_create_collection(self, collection):
                pass

            async def delete_collection(self, collection):
                pass

        return MockVectorDB()

    @pytest.mark.asyncio
    async def test_add_embeddings_signature(self, mock_vector_db):
        """add_embeddings has expected signature."""
        await mock_vector_db.add_embeddings(
            collection='test',
            ids=['id1', 'id2'],
            embeddings_list=[[0.1, 0.2], [0.3, 0.4]],
            metadatas=[{'a': 1}, {'b': 2}],
            documents=['doc1', 'doc2'],
        )
        mock_vector_db.add_embeddings.assert_called_once()

    @pytest.mark.asyncio
    async def test_search_signature(self, mock_vector_db):
        """search has expected signature with all optional params."""
        import numpy as np

        await mock_vector_db.search(
            collection='test',
            query_embedding=np.array([0.1, 0.2]),
            k=10,
            search_type='hybrid',
            query_text='search text',
            filter={'file_id': 'abc'},
            vector_weight=0.7,
        )
        mock_vector_db.search.assert_called_once()

    @pytest.mark.asyncio
    async def test_delete_by_filter_returns_int(self, mock_vector_db):
        """delete_by_filter returns int count."""
        result = await mock_vector_db.delete_by_filter('test', {'file_id': 'abc'})
        assert isinstance(result, int)