1 files changed, 18 insertions, 5 deletions
diff --git a/autogpts/autogpt/tests/mocks/mock_commands.py b/autogpts/autogpt/tests/mocks/mock_commands.py
index 3758c1da2..ab9e961b6 100644
--- a/autogpts/autogpt/tests/mocks/mock_commands.py
+++ b/autogpts/autogpt/tests/mocks/mock_commands.py
@@ -1,16 +1,29 @@
 from autogpt.command_decorator import command
+from autogpt.core.utils.json_schema import JSONSchema
 
 COMMAND_CATEGORY = "mock"
 
 
 @command(
-    "function_based",
+    "function_based_cmd",
     "Function-based test command",
     {
-        "arg1": {"type": "int", "description": "arg 1", "required": True},
-        "arg2": {"type": "str", "description": "arg 2", "required": True},
+        "arg1": JSONSchema(
+            type=JSONSchema.Type.INTEGER,
+            description="arg 1",
+            required=True,
+        ),
+        "arg2": JSONSchema(
+            type=JSONSchema.Type.STRING,
+            description="arg 2",
+            required=True,
+        ),
     },
 )
-def function_based(arg1: int, arg2: str) -> str:
-    """A function-based test command that returns a string with the two arguments separated by a dash."""
+def function_based_cmd(arg1: int, arg2: str) -> str:
+    """A function-based test command.
+
+    Returns:
+        str: the two arguments separated by a dash.
+    """
     return f"{arg1} - {arg2}"