From e51cb768dbfc4a4838e86a5388aea5a41e9d2541 Mon Sep 17 00:00:00 2001
From: Michael Panchenko <michael.panchenko@oraios-ai.de>
Date: Tue, 22 Apr 2025 17:32:45 +0200
Subject: [PATCH] Cleanup after ai

---
 src/serena/agent.py                   |  9 ++++-----
 test/{ => serena}/test_edit_marker.py |  0
 test/serena/test_mcp.py               | 20 ++------------------
 3 files changed, 6 insertions(+), 23 deletions(-)
 rename test/{ => serena}/test_edit_marker.py (100%)

diff --git a/src/serena/agent.py b/src/serena/agent.py
index 86e93a9..df7099f 100644
--- a/src/serena/agent.py
+++ b/src/serena/agent.py
@@ -1372,11 +1372,10 @@ class InitialInstructionsTool(Tool):
 
 def iter_tool_classes(same_module_only: bool = True) -> Generator[type[Tool], None, None]:
     """
-    Iterate over all Tool subclasses.
-    
-    Args:
-        same_module_only: If True, only iterate over tools defined in the same module as the Tool class.
-                         If False, iterate over all Tool subclasses.
+    Iterate over Tool subclasses.
+
+    :param same_module_only: Whether to only iterate over tools defined in the same module as the Tool class
+        or over all subclasses of Tool.
     """
     for tool_class in iter_subclasses(Tool):
         if same_module_only and tool_class.__module__ != Tool.__module__:
diff --git a/test/test_edit_marker.py b/test/serena/test_edit_marker.py
similarity index 100%
rename from test/test_edit_marker.py
rename to test/serena/test_edit_marker.py
diff --git a/test/serena/test_mcp.py b/test/serena/test_mcp.py
index bac38da..b199917 100644
--- a/test/serena/test_mcp.py
+++ b/test/serena/test_mcp.py
@@ -3,7 +3,7 @@
 import pytest
 from mcp.server.fastmcp.tools.base import Tool as MCPTool
 
-from serena.agent import Tool
+from serena.agent import Tool, iter_tool_classes
 from serena.mcp import make_tool
 
 
@@ -269,23 +269,7 @@ def is_test_mock_class(tool_class: type) -> bool:
     )
 
 
-def get_real_tool_classes():
-    """Get all non-test, non-abstract tool classes."""
-    from serena.agent import iter_tool_classes
-
-    for tool_class in iter_tool_classes(same_module_only=False):
-        # Skip abstract base classes that can't be instantiated
-        if tool_class.__name__ == "Tool" or getattr(tool_class, "__abstractmethods__", set()):
-            continue
-
-        # Skip test mock classes
-        if is_test_mock_class(tool_class):
-            continue
-
-        yield tool_class
-
-
-@pytest.mark.parametrize("tool_class", list(get_real_tool_classes()))
+@pytest.mark.parametrize("tool_class", list(iter_tool_classes()))
 def test_make_tool_all_tools(tool_class) -> None:
     """Test that make_tool works for all tools in the codebase."""