From 142b04fa760490062f8366b836784ee02206e491 Mon Sep 17 00:00:00 2001
From: lidanyang <lidanyang@fuzhi.ai>
Date: Fri, 24 Nov 2023 15:04:14 +0800
Subject: [PATCH] test tool register

---
 tests/metagpt/tools/functions/__init__.py     |  6 ++
 .../tools/functions/register/__init__.py      |  6 ++
 .../tools/functions/register/test_register.py | 55 +++++++++++++++++++
 3 files changed, 67 insertions(+)
 create mode 100644 tests/metagpt/tools/functions/__init__.py
 create mode 100644 tests/metagpt/tools/functions/register/__init__.py
 create mode 100644 tests/metagpt/tools/functions/register/test_register.py

diff --git a/tests/metagpt/tools/functions/__init__.py b/tests/metagpt/tools/functions/__init__.py
new file mode 100644
index 000000000..7d36f3404
--- /dev/null
+++ b/tests/metagpt/tools/functions/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2023/11/17 10:24
+# @Author  : lidanyang
+# @File    : __init__.py
+# @Desc    :
diff --git a/tests/metagpt/tools/functions/register/__init__.py b/tests/metagpt/tools/functions/register/__init__.py
new file mode 100644
index 000000000..7d36f3404
--- /dev/null
+++ b/tests/metagpt/tools/functions/register/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2023/11/17 10:24
+# @Author  : lidanyang
+# @File    : __init__.py
+# @Desc    :
diff --git a/tests/metagpt/tools/functions/register/test_register.py b/tests/metagpt/tools/functions/register/test_register.py
new file mode 100644
index 000000000..a71f7d01c
--- /dev/null
+++ b/tests/metagpt/tools/functions/register/test_register.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2023/11/17 10:24
+# @Author  : lidanyang
+# @File    : test_register.py
+# @Desc    :
+import pytest
+
+from metagpt.tools.functions.register.register import FunctionRegistry
+from metagpt.tools.functions.schemas.base import ToolSchema, field
+
+
+@pytest.fixture
+def registry():
+    return FunctionRegistry()
+
+
+class AddNumbers(ToolSchema):
+    """Add two numbers"""
+
+    num1: int = field(description="First number")
+    num2: int = field(description="Second number")
+
+
+def test_register(registry):
+    @registry.register("module1", AddNumbers)
+    def add_numbers(num1, num2):
+        return num1 + num2
+
+    assert len(registry.functions["module1"]) == 1
+    assert "add_numbers" in registry.functions["module1"]
+
+    with pytest.raises(ValueError):
+
+        @registry.register("module1", AddNumbers)
+        def add_numbers(num1, num2):
+            return num1 + num2
+
+    func = registry.get("module1", "add_numbers")
+    assert func["func"](1, 2) == 3
+    assert func["schema"] == {
+        "name": "add_numbers",
+        "description": "Add two numbers",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "num1": {"description": "First number", "type": "int"},
+                "num2": {"description": "Second number", "type": "int"},
+            },
+            "required": ["num1", "num2"],
+        },
+    }
+
+    module1_funcs = registry.get_all_by_module("module1")
+    assert len(module1_funcs) == 1