From 86dae813cd5692715e6bcd2c73e38e877d09d96a Mon Sep 17 00:00:00 2001
From: Kumaran Rajendhiran <kumaran@airt.ai>
Date: Thu, 23 Jan 2025 06:02:48 +0000
Subject: [PATCH] Add neo4j test markers

---
 .../contrib/graph_rag/test_native_neo4j_graph_rag.py          | 3 +++
 test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py      | 4 ++++
 test/agentchat/contrib/test_llamaindex_conversable_agent.py   | 1 +
 3 files changed, 8 insertions(+)

diff --git a/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py b/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py
index a303578a26..29dd38ce3c 100644
--- a/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py
+++ b/test/agentchat/contrib/graph_rag/test_native_neo4j_graph_rag.py
@@ -114,6 +114,7 @@ def neo4j_native_query_engine_auto():
 
 
 @pytest.mark.openai
+@pytest.mark.neo4j
 @pytest.mark.skipif(
     sys.platform in ["darwin", "win32"] or skip,
     reason=reason,
@@ -128,6 +129,7 @@ def test_neo4j_native_query_engine(neo4j_native_query_engine):
 
 
 @pytest.mark.openai
+@pytest.mark.neo4j
 @pytest.mark.skipif(
     sys.platform in ["darwin", "win32"] or skip,
     reason=reason,
@@ -141,6 +143,7 @@ def test_neo4j_native_query_auto(neo4j_native_query_engine_auto):
     assert query_result.answer.find("BUZZ") >= 0
 
 
+@pytest.mark.neo4j
 def test_neo4j_add_records(neo4j_native_query_engine):
     """Test the add_records functionality of the Neo4j Query Engine."""
     input_path = "./test/agentchat/contrib/graph_rag/the_matrix.txt"
diff --git a/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py b/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py
index 79c6faae0e..38141ec01d 100644
--- a/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py
+++ b/test/agentchat/contrib/graph_rag/test_neo4j_graph_rag.py
@@ -121,6 +121,7 @@ def neo4j_query_engine_auto():
 
 
 @pytest.mark.openai
+@pytest.mark.neo4j
 @pytest.mark.skipif(
     sys.platform in ["darwin", "win32"] or skip,
     reason=reason,
@@ -138,6 +139,7 @@ def test_neo4j_query_engine(neo4j_query_engine):
 
 
 @pytest.mark.openai
+@pytest.mark.neo4j
 @pytest.mark.skipif(
     sys.platform in ["darwin", "win32"] or skip,
     reason=reason,
@@ -160,6 +162,7 @@ def test_neo4j_add_records(neo4j_query_engine):
 
 
 @pytest.mark.openai
+@pytest.mark.neo4j
 @pytest.mark.skipif(
     sys.platform in ["darwin", "win32"] or skip,
     reason=reason,
@@ -174,6 +177,7 @@ def test_neo4j_auto(neo4j_query_engine_auto):
 
 
 @pytest.mark.openai
+@pytest.mark.neo4j
 @pytest.mark.skipif(
     sys.platform in ["darwin", "win32"] or skip,
     reason=reason,
diff --git a/test/agentchat/contrib/test_llamaindex_conversable_agent.py b/test/agentchat/contrib/test_llamaindex_conversable_agent.py
index b20e9fe5d2..5579020fe0 100644
--- a/test/agentchat/contrib/test_llamaindex_conversable_agent.py
+++ b/test/agentchat/contrib/test_llamaindex_conversable_agent.py
@@ -29,6 +29,7 @@
 openai_key = MOCK_OPEN_AI_API_KEY
 
 
+@pytest.mark.neo4j
 @pytest.mark.skipif(skip_for_dependencies, reason=skip_reason)
 @patch("llama_index.core.agent.ReActAgent.chat")
 def test_group_chat_with_llama_index_conversable_agent(chat_mock: MagicMock) -> None: