Add tests for function schema with args, kwargs

jackmpcollins · Feb 14, 2024 · 2184380 · 2184380
1 parent c5a1955
commit 2184380
Showing 1 changed file with 57 additions and 5 deletions.
diff --git a/tests/test_function_schema.py b/tests/test_function_schema.py
@@ -530,6 +530,14 @@ def plus_default_value(a: int, b: int = 3) -> int:
     return a + b
 
 
+def plus_with_args(a: int, *args: int) -> int:
+    return a + sum(args)
+
+
+def plus_with_kwargs(a: int, **kwargs: int) -> int:
+    return a + sum(kwargs.values())
+
+
 def plus_with_annotated(
     a: Annotated[int, Field(description="First number")],
     b: Annotated[int, Field(description="Second number")],
@@ -600,6 +608,44 @@ def plus_with_basemodel(a: IntModel, b: IntModel) -> IntModel:
                 },
             },
         ),
+        (
+            plus_with_args,
+            {
+                "name": "plus_with_args",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "a": {"title": "A", "type": "integer"},
+                        "args": {
+                            "default": [],
+                            "items": {"type": "integer"},
+                            "title": "Args",
+                            "type": "array",
+                        },
+                    },
+                    "required": ["a"],
+                },
+            },
+        ),
+        (
+            plus_with_kwargs,
+            {
+                "name": "plus_with_kwargs",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "a": {"title": "A", "type": "integer"},
+                        "kwargs": {
+                            "additionalProperties": {"type": "integer"},
+                            "default": {},
+                            "title": "Kwargs",
+                            "type": "object",
+                        },
+                    },
+                    "required": ["a"],
+                },
+            },
+        ),
         (
             plus_with_annotated,
             {
@@ -662,22 +708,28 @@ def test_function_call_function_schema_with_default_value():
 
 
 function_call_function_schema_args_test_cases = [
-    (plus, '{"a": 1, "b": 2}', FunctionCall(plus, a=1, b=2)),
+    (plus, '{"a": 1, "b": 2}', FunctionCall(plus, 1, 2)),
     (
         plus_no_type_hints,
         '{"a": 1, "b": 2}',
-        FunctionCall(plus_no_type_hints, a=1, b=2),
+        FunctionCall(plus_no_type_hints, 1, 2),
+    ),
+    (plus_default_value, '{"a": 1}', FunctionCall(plus_default_value, 1)),
+    (plus_with_args, '{"a": 1, "args": [2, 3]}', FunctionCall(plus_with_args, 1, 2, 3)),
+    (
+        plus_with_kwargs,
+        '{"a": 1, "kwargs": {"b": 2, "c": 3}}',
+        FunctionCall(plus_with_kwargs, 1, b=2, c=3),
     ),
-    (plus_default_value, '{"a": 1}', FunctionCall(plus_default_value, a=1)),
     (
         plus_with_annotated,
         '{"a": 1, "b": 2}',
-        FunctionCall(plus_with_annotated, a=1, b=2),
+        FunctionCall(plus_with_annotated, 1, 2),
     ),
     (
         plus_with_basemodel,
         '{"a": {"value": 1}, "b": {"value": 2}}',
-        FunctionCall(plus_with_basemodel, a=IntModel(value=1), b=IntModel(value=2)),
+        FunctionCall(plus_with_basemodel, IntModel(value=1), IntModel(value=2)),
     ),
 ]