Cleaned up closing of graphstore and clients in tests.

Author: caseyclements

libs/langchain-mongodb/tests/integration_tests/test_graph_large.py

@@ -13,7 +13,9 @@
 from typing import Generator
 
 import pytest
+from langchain_core.documents import Document
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from pymongo import MongoClient
@@ -85,10 +87,8 @@ def graph_store(
     collection: Collection,
     entity_extraction_model: BaseChatModel,
     text_chunks: list[str],
-) -> MongoDBGraphStore:
+) -> Generator[MongoDBGraphStore, None, None]:
     """Create MongoDBGraphStore and populate it with entities from text chunks."""
-    from langchain_core.documents import Document
-
     # Create the graph store
     store = MongoDBGraphStore(
         collection=collection,
@@ -105,7 +105,8 @@ def graph_store(
     # Verify that we processed all chunks
     assert len(bulkwrite_results) == len(documents)
 
-    return store
+    yield store
+    store.close()
 
 
 def test_large_graph_entity_count(graph_store: MongoDBGraphStore):
@@ -257,8 +258,6 @@ def test_chunk_count(text_chunks: list[str]):
 
 def test_graph_store_chat_response(graph_store: MongoDBGraphStore):
     """Test the chat response functionality with the large graph."""
-    from langchain_core.messages import AIMessage
-
     query = (
         "Tell me about the AI Ethics Initiative and which companies are supporting it."
     )

libs/langchain-mongodb/tests/integration_tests/test_graphrag.py

@@ -201,13 +201,27 @@ def test_additional_entity_examples(entity_extraction_model, entity_example, doc
         assert len(new_entities) >= 2
 
 
+def test_chat_response(graph_store, query_connection):
+    """Displays querying an existing Knowledge Graph Database"""
+    answer = graph_store.chat_response(query_connection)
+    assert isinstance(answer, AIMessage)
+    assert "acme corporation" in answer.content.lower()
+
+
+def test_similarity_search(graph_store, query_connection):
+    docs = graph_store.similarity_search(query_connection)
+    assert len(docs) >= 4
+    assert all({"_id", "type", "relationships"}.issubset(set(d.keys())) for d in docs)
+    assert any("depth" in d.keys() for d in docs)
+    assert any("attributes" in d.keys() for d in docs)
+
+
 def test_validator(documents, entity_extraction_model):
     # Case 1. No existing collection.
+    client = MongoClient(CONNECTION_STRING)
     clxn_name = f"{COLLECTION_NAME}_validation"
-
-    with MongoClient(CONNECTION_STRING) as client:
-        client[DB_NAME][clxn_name].drop()
-
+    client[DB_NAME][clxn_name].drop()
+    client.close()
     # now we call with validation that can be added without db admin privileges
     store = MongoDBGraphStore(
         connection_string=CONNECTION_STRING,
@@ -229,15 +243,16 @@ def test_validator(documents, entity_extraction_model):
         collection = client[DB_NAME][clxn_name]
         collection.delete_many({})
 
-        store = MongoDBGraphStore(
-            collection=collection,
-            entity_extraction_model=entity_extraction_model,
-            validate=True,
-            validation_action="error",
-        )
-        bulkwrite_results = store.add_documents(documents)
-        assert len(bulkwrite_results) == len(documents)
-        collection.delete_many({})
+    store = MongoDBGraphStore(
+        collection=collection,
+        entity_extraction_model=entity_extraction_model,
+        validate=True,
+        validation_action="error",
+    )
+    bulkwrite_results = store.add_documents(documents)
+    assert len(bulkwrite_results) == len(documents)
+    collection.drop()
+    store.close()
 
     # Case 3: Existing collection without a validator
     with MongoClient(CONNECTION_STRING) as client: