Skip to content

Commit

Permalink
Skip test if mongodb is unreachable
Browse files Browse the repository at this point in the history
  • Loading branch information
@kumaranvpl
kumaranvpl committed Jan 24, 2025
1 parent a3a7b98 commit 46a3976
Showing 1 changed file with 55 additions and 0 deletions.
55 changes: 55 additions & 0 deletions test/agentchat/contrib/vectordb/test_mongodb.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import logging
import os
import random
import sys
from time import monotonic, sleep

import pytest
Expand All @@ -31,6 +32,18 @@
DELAY = 2
TIMEOUT = 120.0

reason = "do not run on MacOS or windows OR if mongodb is not running"


def is_mongodb_accessible():
try:
client = MongoClient(MONGODB_URI, serverSelectionTimeoutMS=2000)
client.admin.command("ping")
client.close()
return True
except Exception:
return False


def _wait_for_predicate(predicate, err, timeout=TIMEOUT, interval=DELAY):
"""Generic to block until the predicate returns true
Expand Down Expand Up @@ -133,6 +146,9 @@ def collection_name():
return f"{MONGODB_COLLECTION}_{collection_id}"


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_create_collection(db, collection_name):
"""Def create_collection(collection_name: str,
Expand Down Expand Up @@ -163,6 +179,9 @@ def test_create_collection(db, collection_name):
db.create_collection(collection_name=collection_name, overwrite=False, get_or_create=False)


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_get_collection(db, collection_name):
with pytest.raises(ValueError):
Expand All @@ -177,6 +196,9 @@ def test_get_collection(db, collection_name):
assert collection_got.name == db.active_collection.name


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_delete_collection(db, collection_name):
assert collection_name not in db.list_collections()
Expand All @@ -186,6 +208,9 @@ def test_delete_collection(db, collection_name):
assert collection_name not in db.list_collections()


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_insert_docs(db, collection_name, example_documents):
# Test that there's an active collection
Expand All @@ -212,6 +237,9 @@ def test_insert_docs(db, collection_name, example_documents):
assert len(found[0]["embedding"]) == 384


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_update_docs(db_with_indexed_clxn, example_documents):
db, collection = db_with_indexed_clxn
Expand Down Expand Up @@ -248,6 +276,9 @@ def test_update_docs(db_with_indexed_clxn, example_documents):
assert collection.find_one({"_id": new_id}) is None


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_delete_docs(db_with_indexed_clxn, example_documents):
db, clxn = db_with_indexed_clxn
Expand All @@ -259,6 +290,9 @@ def test_delete_docs(db_with_indexed_clxn, example_documents):
assert {2, "2"} == {doc["_id"] for doc in clxn.find({})}


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_get_docs_by_ids(db_with_indexed_clxn, example_documents):
db, clxn = db_with_indexed_clxn
Expand All @@ -285,12 +319,18 @@ def test_get_docs_by_ids(db_with_indexed_clxn, example_documents):
assert len(docs) == 4


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_retrieve_docs_empty(db_with_indexed_clxn):
db, clxn = db_with_indexed_clxn
assert db.retrieve_docs(queries=["Cats"], collection_name=clxn.name, n_results=2) == []


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_retrieve_docs_populated_db_empty_query(db_with_indexed_clxn, example_documents):
db, clxn = db_with_indexed_clxn
Expand All @@ -300,6 +340,9 @@ def test_retrieve_docs_populated_db_empty_query(db_with_indexed_clxn, example_do
assert results == []


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_retrieve_docs(db_with_indexed_clxn, example_documents):
"""Begin testing Atlas Vector Search
Expand All @@ -324,6 +367,9 @@ def results_ready():
assert all(["embedding" not in doc[0] for doc in results[0]])


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_retrieve_docs_with_embedding(db_with_indexed_clxn, example_documents):
"""Begin testing Atlas Vector Search
Expand All @@ -348,6 +394,9 @@ def results_ready():
assert all(["embedding" in doc[0] for doc in results[0]])


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_retrieve_docs_multiple_queries(db_with_indexed_clxn, example_documents):
db, clxn = db_with_indexed_clxn
Expand All @@ -371,6 +420,9 @@ def results_ready():
assert {doc[0]["id"] for doc in results[1]} == {"1", "2"}


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_retrieve_docs_with_threshold(db_with_indexed_clxn, example_documents):
db, clxn = db_with_indexed_clxn
Expand All @@ -393,6 +445,9 @@ def results_ready():
assert all([doc[1] >= 0.7 for doc in results[0]])


@pytest.mark.skipif(
sys.platform in ["darwin", "win32"] or not is_mongodb_accessible(),
)
@skip_on_missing_imports(["pymongo", "sentence_transformers"], "retrievechat-mongodb")
def test_wait_until_document_ready(collection_name, example_documents):
database = MongoClient(MONGODB_URI)[MONGODB_DATABASE]
Expand Down

0 comments on commit 46a3976

Please sign in to comment.