Add first tests

Author: NoahStapp

.gitignore

@@ -1,2 +1,3 @@
 .idea/
 .DS_STORE
+__pycache__/

CHANGELOG.md

@@ -4,4 +4,4 @@ ______________________________________________________________________
 
 ## Changes in version 0.1 (2025/09/XX)
 
-- Initial release.
+- Initial release.

pymongo_vectorsearch_utils/__init__.py

@@ -0,0 +1,27 @@
+from ._version import __version__
+from .index import (
+    create_fulltext_search_index,
+    create_vector_search_index,
+    drop_vector_search_index,
+    update_vector_search_index,
+)
+from .pipeline import (
+    combine_pipelines,
+    final_hybrid_stage,
+    reciprocal_rank_stage,
+    text_search_stage,
+    vector_search_stage,
+)
+
+__all__ = [
+    "__version__",
+    "create_vector_search_index",
+    "drop_vector_search_index",
+    "update_vector_search_index",
+    "create_fulltext_search_index",
+    "text_search_stage",
+    "vector_search_stage",
+    "combine_pipelines",
+    "reciprocal_rank_stage",
+    "final_hybrid_stage",
+]

pymongo_vectorsearch_utils/index.py

@@ -186,6 +186,31 @@ def update_vector_search_index(
     logger.info("Update succeeded")
 
 
+def drop_vector_search_index(
+    collection: Collection[Any],
+    index_name: str,
+    *,
+    wait_until_complete: float | None = None,
+) -> None:
+    """Drop a created vector search index.
+
+    Args:
+        collection (Collection): MongoDB Collection with index to be dropped.
+        index_name (str): Name of the MongoDB index.
+        wait_until_complete (Optional[float]): If provided, number of seconds to wait
+            until search index is ready.
+    """
+    logger.info("Dropping Search Index %s from Collection: %s", index_name, collection.name)
+    collection.drop_search_index(index_name)
+    if wait_until_complete:
+        wait_for_predicate(
+            predicate=lambda: len(list(collection.list_search_indexes())) == 0,
+            err=f"Index {index_name} did not drop in {wait_until_complete}!",
+            timeout=wait_until_complete,
+        )
+    logger.info("Vector Search index %s.%s dropped", collection.name, index_name)
+
+
 def create_fulltext_search_index(
     collection: Collection[Any],
     index_name: str,

pyproject.toml

@@ -62,6 +62,7 @@ testpaths = [
 python_version = "3.10"
 strict = true
 enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
+exclude = ["tests/"]
 
 [tool.ruff]
 line-length = 100

tests/__init__.py

@@ -0,0 +1,114 @@
+import os
+from collections.abc import Generator
+
+import pytest
+from pymongo import MongoClient
+from pymongo.collection import Collection
+
+from pymongo_vectorsearch_utils.index import (
+    create_vector_search_index,
+    drop_vector_search_index,
+    is_index_ready,
+    update_vector_search_index,
+)
+
+DB_NAME = "vectorsearch_utils_test"
+COLLECTION_NAME = "test_index"
+VECTOR_INDEX_NAME = "vector_index"
+
+TIMEOUT = 120
+DIMENSIONS = 10
+
+
+@pytest.fixture(scope="module")
+def client() -> Generator[MongoClient, None, None]:
+    conn_str = os.environ.get("MONGODB_URI", "mongodb://127.0.0.1:27017?directConnection=true")
+    client = MongoClient(conn_str)
+    yield client
+    client.close()
+
+
+@pytest.fixture
+def collection(client) -> Generator:
+    if COLLECTION_NAME not in client[DB_NAME].list_collection_names():
+        clxn = client[DB_NAME].create_collection(COLLECTION_NAME)
+    else:
+        clxn = client[DB_NAME][COLLECTION_NAME]
+    clxn = client[DB_NAME][COLLECTION_NAME]
+    clxn.delete_many({})
+    yield clxn
+    clxn.delete_many({})
+
+
+def test_search_index_create_and_drop(collection: Collection) -> None:
+    index_name = VECTOR_INDEX_NAME
+    dimensions = DIMENSIONS
+    path = "embedding"
+    similarity = "cosine"
+    filters: list[str] | None = None
+    wait_until_complete = TIMEOUT
+
+    for index_info in collection.list_search_indexes():
+        drop_vector_search_index(
+            collection, index_info["name"], wait_until_complete=wait_until_complete
+        )
+
+    assert len(list(collection.list_search_indexes())) == 0
+
+    create_vector_search_index(
+        collection=collection,
+        index_name=index_name,
+        dimensions=dimensions,
+        path=path,
+        similarity=similarity,
+        filters=filters,
+        wait_until_complete=wait_until_complete,
+    )
+
+    assert is_index_ready(collection, index_name)
+    indexes = list(collection.list_search_indexes())
+    assert len(indexes) == 1
+    assert indexes[0]["name"] == index_name
+
+    drop_vector_search_index(collection, index_name, wait_until_complete=wait_until_complete)
+
+    indexes = list(collection.list_search_indexes())
+    assert len(indexes) == 0
+
+
+def test_search_index_update_vector_search_index(collection: Collection) -> None:
+    index_name = "INDEX_TO_UPDATE"
+    similarity_orig = "cosine"
+    similarity_new = "euclidean"
+
+    # Create another index
+    create_vector_search_index(
+        collection=collection,
+        index_name=index_name,
+        dimensions=DIMENSIONS,
+        path="embedding",
+        similarity=similarity_orig,
+        wait_until_complete=TIMEOUT,
+    )
+
+    assert is_index_ready(collection, index_name)
+    indexes = list(collection.list_search_indexes())
+    assert len(indexes) == 1
+    assert indexes[0]["name"] == index_name
+    assert indexes[0]["latestDefinition"]["fields"][0]["similarity"] == similarity_orig
+
+    # Update the index and test new similarity
+    update_vector_search_index(
+        collection=collection,
+        index_name=index_name,
+        dimensions=DIMENSIONS,
+        path="embedding",
+        similarity=similarity_new,
+        wait_until_complete=TIMEOUT,
+    )
+
+    assert is_index_ready(collection, index_name)
+    indexes = list(collection.list_search_indexes())
+    assert len(indexes) == 1
+    assert indexes[0]["name"] == index_name
+    assert indexes[0]["latestDefinition"]["fields"][0]["similarity"] == similarity_new