remove ToolTypesEnum

tests/metagpt/roles/test_ml_engineer.py

@@ -4,7 +4,7 @@ from metagpt.actions.execute_nb_code import ExecuteNbCode
 from metagpt.logs import logger
 from metagpt.roles.ml_engineer import MLEngineer
 from metagpt.schema import Message, Plan, Task
-from metagpt.tools.tool_data_type import ToolTypeEnum
+from metagpt.tools.tool_types import ToolTypes
 from tests.metagpt.actions.test_debug_code import CODE, DebugContext, ErrorStr
 
 
@@ -63,7 +63,7 @@ async def test_mle_update_data_columns(mocker):
     mle.planner.plan = MockPlan
 
     # manually update task type to test update
-    mle.planner.plan.current_task.task_type = ToolTypeEnum.DATA_PREPROCESS.value
+    mle.planner.plan.current_task.task_type = ToolTypes.DATA_PREPROCESS.value
 
     result = await mle._update_data_columns()
     assert result is not None

tests/metagpt/tools/test_tool_registry.py

@@ -1,7 +1,7 @@
 import pytest
 
 from metagpt.tools.tool_registry import ToolRegistry
-from metagpt.tools.tool_types import ToolType
+from metagpt.tools.tool_types import ToolTypes
 
 
 @pytest.fixture
@@ -9,6 +9,11 @@ def tool_registry():
     return ToolRegistry()
 
 
+@pytest.fixture
+def tool_registry_full():
+    return ToolRegistry(tool_types=ToolTypes)
+
+
 @pytest.fixture
 def schema_yaml(mocker):
     mock_yaml_content = """
@@ -29,11 +34,12 @@ def test_initialization(tool_registry):
     assert tool_registry.tools_by_types == {}
 
 
-# Test Tool Type Registration
-def test_register_tool_type(tool_registry):
-    tool_type = ToolType(name="TestType", desc="test")
-    tool_registry.register_tool_type(tool_type)
-    assert "TestType" in tool_registry.tool_types
+# Test Initialization with tool types
+def test_initialize_with_tool_types(tool_registry_full):
+    assert isinstance(tool_registry_full, ToolRegistry)
+    assert tool_registry_full.tools == {}
+    assert tool_registry_full.tools_by_types == {}
+    assert "data_preprocess" in tool_registry_full.tool_types
 
 
 # Test Tool Registration
@@ -66,27 +72,21 @@ def test_get_tool(tool_registry, schema_yaml):
 
 
 # Similar tests for has_tool_type, get_tool_type, get_tools_by_type
-def test_has_tool_type(tool_registry):
-    tool_type = ToolType(name="TestType", desc="test")
-    tool_registry.register_tool_type(tool_type)
-    assert tool_registry.has_tool_type("TestType")
-    assert not tool_registry.has_tool_type("NonexistentType")
+def test_has_tool_type(tool_registry_full):
+    assert tool_registry_full.has_tool_type("data_preprocess")
+    assert not tool_registry_full.has_tool_type("NonexistentType")
 
 
-def test_get_tool_type(tool_registry):
-    tool_type = ToolType(name="TestType", desc="test")
-    tool_registry.register_tool_type(tool_type)
-    retrieved_type = tool_registry.get_tool_type("TestType")
+def test_get_tool_type(tool_registry_full):
+    retrieved_type = tool_registry_full.get_tool_type("data_preprocess")
     assert retrieved_type is not None
-    assert retrieved_type.name == "TestType"
+    assert retrieved_type.name == "data_preprocess"
 
 
 def test_get_tools_by_type(tool_registry, schema_yaml):
     tool_type_name = "TestType"
     tool_name = "TestTool"
     tool_path = "/path/to/tool"
-    tool_type = ToolType(name=tool_type_name, desc="test")
-    tool_registry.register_tool_type(tool_type)
 
     tool_registry.register_tool(tool_name, tool_path, tool_type=tool_type_name)