INTPYTHON-629 Address flaky failures in test_graphrag (#146)

Author: caseyclements

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

@@ -69,7 +69,7 @@ def documents():
 GreenTech Ltd. Leads SolarGrid Initiative
 
 San Francisco, CA — GreenTech Ltd. has emerged as a leader in renewable energy projects with the SolarGrid Initiative,
-a collaboration with ACME Corporation. Jane Smith, the project’s Lead Manager, highlighted its ambitious goal:
+a collaboration with ACME Corporation. Jane Smith, the SolarGrid project’s Lead Manager, highlighted its ambitious goal:
 providing affordable solar energy to underserved communities.
 
 GreenTech, headquartered in San Francisco, has worked closely with ACME since their partnership began in May 2021.
@@ -136,7 +136,7 @@ def graph_store(collection, entity_extraction_model, documents) -> MongoDBGraphS
 
 @pytest.fixture
 def query_connection():
-    return "How are Jane Smith and John Doe related?"
+    return "How are the project manager Jane Smith and the Chief Technology Officer John Doe related?"
 
 
 def test_add_docs_store(graph_store):
@@ -145,17 +145,19 @@ def test_add_docs_store(graph_store):
     assert 4 <= len(extracted_entities) < 8
 
 
-@flaky
+@flaky(max_runs=3, min_passes=2)
 def test_extract_entity_names(graph_store, query_connection):
     query_entity_names = graph_store.extract_entity_names(query_connection)
     assert set(query_entity_names) == {"John Doe", "Jane Smith"}
 
+
+def test_extract_entities_from_empty_string_names(graph_store):
     no_names = graph_store.extract_entity_names("")
     assert isinstance(no_names, list)
     assert len(no_names) == 0
 
 
-@flaky
+@flaky(max_runs=3, min_passes=1)
 def test_related_entities(graph_store):
     entity_names = ["John Doe", "Jane Smith"]
     related_entities = graph_store.related_entities(entity_names)
@@ -166,7 +168,7 @@ def test_related_entities(graph_store):
     assert len(no_entities) == 0
 
 
-@flaky
+@flaky(max_runs=3, min_passes=1)
 def test_additional_entity_examples(entity_extraction_model, entity_example, documents):
     # First, create one client just to drop any existing collections
     client = MongoClient(CONNECTION_STRING)
@@ -190,15 +192,15 @@ def test_additional_entity_examples(entity_extraction_model, entity_example, doc
     assert len(new_entities) >= 2
 
 
-@flaky
+@flaky(max_runs=3, min_passes=1)
 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()
 
 
-@flaky
+@flaky(max_runs=3, min_passes=1)
 def test_similarity_search(graph_store, query_connection):
     docs = graph_store.similarity_search(query_connection)
     assert len(docs) >= 4
@@ -207,7 +209,7 @@ def test_similarity_search(graph_store, query_connection):
     assert any("attributes" in d.keys() for d in docs)
 
 
-@flaky
+@flaky(max_runs=3, min_passes=1)
 def test_validator(documents, entity_extraction_model):
     # Case 1. No existing collection.
     client = MongoClient(CONNECTION_STRING)
@@ -261,7 +263,7 @@ def test_validator(documents, entity_extraction_model):
     client.close()
 
 
-@flaky
+@flaky(max_runs=3, min_passes=1)
 def test_allowed_entity_types(documents, entity_extraction_model):
     """Add allowed_entity_types. Use the validator to confirm behaviour."""
     allowed_entity_types = ["Person"]
@@ -289,7 +291,7 @@ def test_allowed_entity_types(documents, entity_extraction_model):
     all([len(e["relationships"].get("attributes", [])) == 0 for e in entities])
 
 
-@flaky
+@flaky(max_runs=3, min_passes=1)
 def test_allowed_relationship_types(documents, entity_extraction_model):
     # drop collection
     client = MongoClient(CONNECTION_STRING)