diff --git a/test/agentchat/contrib/vectordb/test_mongodb.py b/test/agentchat/contrib/vectordb/test_mongodb.py index 9464392a1..e14e484d1 100644 --- a/test/agentchat/contrib/vectordb/test_mongodb.py +++ b/test/agentchat/contrib/vectordb/test_mongodb.py @@ -7,6 +7,7 @@ import logging import os import random +import sys from time import monotonic, sleep import pytest @@ -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 @@ -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, @@ -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): @@ -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() @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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 @@ -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]