Refactor test cases to simplify async context management and improve readability

pytest.ini

@@ -3,4 +3,5 @@ testpaths = tests
 markers =
   slow: slow tests
   git: tests that require git
-addopts = --ignore=prometheus/script/
+addopts = --ignore=prometheus/script/
+asyncio_mode = auto

tests/app/services/test_issue_service.py

@@ -44,7 +44,6 @@ def issue_service(mock_neo4j_service, mock_llm_service, mock_repository_service)
     )
 
 
-@pytest.mark.asyncio
 async def test_answer_issue_with_general_container(issue_service, monkeypatch):
     # Setup
     mock_issue_graph = Mock()
@@ -106,7 +105,6 @@ async def test_answer_issue_with_general_container(issue_service, monkeypatch):
     assert result == ("test_patch", True, True, True, "test_response", IssueType.BUG)
 
 
-@pytest.mark.asyncio
 async def test_answer_issue_with_user_defined_container(issue_service, monkeypatch):
     # Setup
     mock_issue_graph = Mock()

tests/app/services/test_knowledge_graph_service.py

@@ -41,7 +41,6 @@ def knowledge_graph_service(mock_neo4j_service, mock_kg_handler):
     return knowledge_graph_service
 
 
-@pytest.mark.asyncio
 async def test_build_and_save_knowledge_graph(knowledge_graph_service, mock_kg_handler):
     """Test the build_and_save_knowledge_graph method."""
     # Given

tests/app/services/test_neo4j_service.py

@@ -6,12 +6,12 @@
 
 @pytest.mark.slow
 async def test_neo4j_service(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        neo4j_service = Neo4jService(
-            neo4j_container.get_connection_url(), neo4j_container.username, neo4j_container.password
-        )
-        assert neo4j_service.neo4j_driver is not None
-        try:
-            neo4j_service.neo4j_driver.verify_connectivity()
-        except Exception as e:
-            pytest.fail(f"Connection verification failed: {e}")
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    neo4j_service = Neo4jService(
+        neo4j_container.get_connection_url(), neo4j_container.username, neo4j_container.password
+    )
+    assert neo4j_service.neo4j_driver is not None
+    try:
+        neo4j_service.neo4j_driver.verify_connectivity()
+    except Exception as e:
+        pytest.fail(f"Connection verification failed: {e}")

tests/graph/test_knowledge_graph.py

@@ -74,9 +74,9 @@ async def test_get_file_tree_depth_two_max_seven_lines():
 
 @pytest.mark.slow
 async def test_from_neo4j(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        driver = neo4j_container.get_driver()
-        handler = KnowledgeGraphHandler(driver, 100)
-        read_kg = handler.read_knowledge_graph(0, 1000, 100, 10)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    driver = neo4j_container.get_driver()
+    handler = KnowledgeGraphHandler(driver, 100)
+    read_kg = handler.read_knowledge_graph(0, 1000, 100, 10)
 
-        assert read_kg == kg
+    assert read_kg == kg

tests/lang_graph/nodes/test_context_provider_node.py

@@ -9,29 +9,27 @@
 @pytest.mark.slow
 async def test_context_provider_node_basic_query(neo4j_container_with_kg_fixture):  # noqa: F811
     """Test basic query handling with the ContextProviderNode."""
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        fake_response = "Fake response"
-        fake_llm = FakeListChatWithToolsModel(responses=[fake_response])
-        node = ContextProviderNode(
-            model=fake_llm,
-            kg=kg,
-            neo4j_driver=neo4j_container.get_driver(),
-            max_token_per_result=1000,
-        )
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    fake_response = "Fake response"
+    fake_llm = FakeListChatWithToolsModel(responses=[fake_response])
+    node = ContextProviderNode(
+        model=fake_llm,
+        kg=kg,
+        neo4j_driver=neo4j_container.get_driver(),
+        max_token_per_result=1000,
+    )
 
-        test_messages = [
-            AIMessage(content="This code handles file processing"),
-            ToolMessage(
-                content="Found implementation in utils.py", tool_call_id="test_tool_call_1"
-            ),
-        ]
-        test_state = {
-            "original_query": "How does the error handling work?",
-            "context_provider_messages": test_messages,
-        }
+    test_messages = [
+        AIMessage(content="This code handles file processing"),
+        ToolMessage(content="Found implementation in utils.py", tool_call_id="test_tool_call_1"),
+    ]
+    test_state = {
+        "original_query": "How does the error handling work?",
+        "context_provider_messages": test_messages,
+    }
 
-        result = node(test_state)
+    result = node(test_state)
 
-        assert "context_provider_messages" in result
-        assert len(result["context_provider_messages"]) == 1
-        assert result["context_provider_messages"][0].content == fake_response
+    assert "context_provider_messages" in result
+    assert len(result["context_provider_messages"]) == 1
+    assert result["context_provider_messages"][0].content == fake_response

tests/neo4j/test_knowledge_graph_handler.py

@@ -11,109 +11,102 @@
 
 @pytest.mark.slow
 async def test_num_ast_nodes(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_ast_nodes = session.execute_read(handler._read_ast_nodes, root_node_id=0)
-                assert len(read_ast_nodes) == 84
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_ast_nodes = session.execute_read(handler._read_ast_nodes, root_node_id=0)
+            assert len(read_ast_nodes) == 84
 
 
 @pytest.mark.slow
 async def test_num_file_nodes(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_file_nodes = session.execute_read(handler._read_file_nodes, root_node_id=0)
-                assert len(read_file_nodes) == 9
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_file_nodes = session.execute_read(handler._read_file_nodes, root_node_id=0)
+            assert len(read_file_nodes) == 9
 
 
 @pytest.mark.slow
 async def test_num_text_nodes(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_text_nodes = session.execute_read(handler._read_text_nodes, root_node_id=0)
-                assert len(read_text_nodes) == 2
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_text_nodes = session.execute_read(handler._read_text_nodes, root_node_id=0)
+            assert len(read_text_nodes) == 2
 
 
 @pytest.mark.slow
 async def test_num_parent_of_edges(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_parent_of_edges = session.execute_read(
-                    handler._read_parent_of_edges, root_node_id=0
-                )
-                assert len(read_parent_of_edges) == 81
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_parent_of_edges = session.execute_read(
+                handler._read_parent_of_edges, root_node_id=0
+            )
+            assert len(read_parent_of_edges) == 81
 
 
 @pytest.mark.slow
 async def test_num_has_file_edges(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_has_file_edges = session.execute_read(
-                    handler._read_has_file_edges, root_node_id=0
-                )
-                print(read_has_file_edges)
-                assert len(read_has_file_edges) == 8
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_has_file_edges = session.execute_read(handler._read_has_file_edges, root_node_id=0)
+            assert len(read_has_file_edges) == 8
 
 
 @pytest.mark.slow
 async def test_num_has_ast_edges(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_has_ast_edges = session.execute_read(
-                    handler._read_has_ast_edges, root_node_id=0
-                )
-                assert len(read_has_ast_edges) == 3
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_has_ast_edges = session.execute_read(handler._read_has_ast_edges, root_node_id=0)
+            assert len(read_has_ast_edges) == 3
 
 
 @pytest.mark.slow
 async def test_num_has_text_edges(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_has_text_edges = session.execute_read(
-                    handler._read_has_text_edges, root_node_id=0
-                )
-                assert len(read_has_text_edges) == 2
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_has_text_edges = session.execute_read(handler._read_has_text_edges, root_node_id=0)
+            assert len(read_has_text_edges) == 2
 
 
 @pytest.mark.slow
 async def test_num_next_chunk_edges(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        with neo4j_container.get_driver() as driver:
-            with driver.session() as session:
-                read_next_chunk_edges = session.execute_read(
-                    handler._read_next_chunk_edges, root_node_id=0
-                )
-                assert len(read_next_chunk_edges) == 1
+    with neo4j_container.get_driver() as driver:
+        with driver.session() as session:
+            read_next_chunk_edges = session.execute_read(
+                handler._read_next_chunk_edges, root_node_id=0
+            )
+            assert len(read_next_chunk_edges) == 1
 
 
 @pytest.mark.slow
 async def test_knowledge_graph_exists(neo4j_container_with_kg_fixture):  # noqa: F811
-    async for neo4j_container, kg in neo4j_container_with_kg_fixture:
-        handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
+    neo4j_container, kg = neo4j_container_with_kg_fixture
+    handler = KnowledgeGraphHandler(neo4j_container.get_driver(), 100)
 
-        assert handler.knowledge_graph_exists(0)
+    assert handler.knowledge_graph_exists(0)
 
 
 @pytest.mark.slow