chore: Add unit tests for the tool and client classes

Author: anubhav756

packages/toolbox-core/src/toolbox_core/client.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+
 import types
 from typing import Any, Callable, Mapping, Optional, Union
 

packages/toolbox-core/tests/test_client.py

@@ -19,6 +19,7 @@
 import pytest
 import pytest_asyncio
 from aioresponses import CallbackResult
+from unittest.mock import AsyncMock
 
 from toolbox_core import ToolboxClient
 from toolbox_core.protocol import ManifestSchema, ParameterSchema, ToolSchema
@@ -301,3 +302,99 @@ async def test_bind_param_fail(self, tool_name, client):
 
         with pytest.raises(Exception):
             tool = tool.bind_parameters({"argC": lambda: 5})
+
+
+@pytest.mark.asyncio
+async def test_new_invoke_tool_server_error(aioresponses, test_tool_str):
+    """Tests that invoking a tool raises an Exception when the server returns an
+    error status."""
+    TOOL_NAME = "server_error_tool"
+    ERROR_MESSAGE = "Simulated Server Error"
+    manifest = ManifestSchema(serverVersion="0.0.0", tools={TOOL_NAME: test_tool_str})
+
+    aioresponses.get(
+        f"{TEST_BASE_URL}/api/tool/{TOOL_NAME}",
+        payload=manifest.model_dump(),
+        status=200,
+    )
+    aioresponses.post(
+        f"{TEST_BASE_URL}/api/tool/{TOOL_NAME}/invoke",
+        payload={"error": ERROR_MESSAGE},
+        status=500,
+    )
+
+    async with ToolboxClient(TEST_BASE_URL) as client:
+        loaded_tool = await client.load_tool(TOOL_NAME)
+
+        with pytest.raises(Exception, match=ERROR_MESSAGE):
+            await loaded_tool(param1="some input")
+
+
+@pytest.mark.asyncio
+async def test_bind_param_async_callable_value_success(aioresponses, test_tool_int_bool):
+    """
+    Tests bind_parameters method with an async callable value.
+    """
+    TOOL_NAME = "async_bind_tool"
+    manifest = ManifestSchema(serverVersion="0.0.0", tools={TOOL_NAME: test_tool_int_bool})
+
+    aioresponses.get(
+        f"{TEST_BASE_URL}/api/tool/{TOOL_NAME}",
+        payload=manifest.model_dump(), status=200
+    )
+
+    def reflect_parameters(url, **kwargs):
+        received_params = kwargs.get("json", {})
+        return CallbackResult(status=200, payload={"result": received_params})
+
+    aioresponses.post(
+        f"{TEST_BASE_URL}/api/tool/{TOOL_NAME}/invoke",
+        callback=reflect_parameters,
+    )
+
+    bound_value_result = True
+    bound_async_callable = AsyncMock(return_value=bound_value_result)
+
+    async with ToolboxClient(TEST_BASE_URL) as client:
+        tool = await client.load_tool(TOOL_NAME)
+        bound_tool = tool.bind_parameters({"argB": bound_async_callable})
+
+        assert bound_tool is not tool
+        assert "argB" not in bound_tool.__signature__.parameters
+        assert "argA" in bound_tool.__signature__.parameters
+
+        passed_value_a = 42
+        res_payload = await bound_tool(argA=passed_value_a)
+
+        assert res_payload == {"argA": passed_value_a, "argB": bound_value_result}
+        bound_async_callable.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_new_add_auth_token_getters_duplicate_fail(aioresponses, test_tool_auth):
+    """
+    Tests that adding a duplicate auth token getter raises ValueError.
+    """
+    TOOL_NAME = "duplicate_auth_tool"
+    AUTH_SERVICE = "my-auth-service"
+    manifest = ManifestSchema(serverVersion="0.0.0", tools={TOOL_NAME: test_tool_auth})
+
+    aioresponses.get(
+        f"{TEST_BASE_URL}/api/tool/{TOOL_NAME}",
+        payload=manifest.model_dump(), status=200
+    )
+
+    def token_handler_1():
+        return "token1"
+
+    def token_handler_2():
+        return "token2"
+
+    async with ToolboxClient(TEST_BASE_URL) as client:
+        tool = await client.load_tool(TOOL_NAME)
+
+        authed_tool = tool.add_auth_token_getters({AUTH_SERVICE: token_handler_1})
+        assert AUTH_SERVICE in authed_tool._ToolboxTool__auth_service_token_getters
+
+        with pytest.raises(ValueError, match=f"Authentication source\\(s\\) `{AUTH_SERVICE}` already registered in tool `{TOOL_NAME}`."):
+             authed_tool.add_auth_token_getters({AUTH_SERVICE: token_handler_2})

packages/toolbox-core/tests/test_protocol.py

@@ -0,0 +1,110 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from inspect import Parameter
+
+import pytest
+
+from toolbox_core.protocol import ParameterSchema
+
+
+def test_parameter_schema_float():
+    """Tests ParameterSchema with type 'float'."""
+    schema = ParameterSchema(name="price", type="float", description="The item price")
+    expected_type = float
+    assert schema._ParameterSchema__get_type() == expected_type
+
+    param = schema.to_param()
+    assert isinstance(param, Parameter)
+    assert param.name == "price"
+    assert param.annotation == expected_type
+    assert param.kind == Parameter.POSITIONAL_OR_KEYWORD
+    assert param.default == Parameter.empty
+
+
+def test_parameter_schema_boolean():
+    """Tests ParameterSchema with type 'boolean'."""
+    schema = ParameterSchema(
+        name="is_active", type="boolean", description="Activity status"
+    )
+    expected_type = bool
+    assert schema._ParameterSchema__get_type() == expected_type
+
+    param = schema.to_param()
+    assert isinstance(param, Parameter)
+    assert param.name == "is_active"
+    assert param.annotation == expected_type
+    assert param.kind == Parameter.POSITIONAL_OR_KEYWORD
+
+
+def test_parameter_schema_array_string():
+    """Tests ParameterSchema with type 'array' containing strings."""
+    item_schema = ParameterSchema(
+        name="", type="string", description=""
+    )
+    schema = ParameterSchema(
+        name="tags", type="array", description="List of tags", items=item_schema
+    )
+
+    assert schema._ParameterSchema__get_type() == list[str]
+
+    param = schema.to_param()
+    assert isinstance(param, Parameter)
+    assert param.name == "tags"
+    assert param.annotation == list[str]
+    assert param.kind == Parameter.POSITIONAL_OR_KEYWORD
+
+
+def test_parameter_schema_array_integer():
+    """Tests ParameterSchema with type 'array' containing integers."""
+    item_schema = ParameterSchema(name="", type="integer", description="")
+    schema = ParameterSchema(
+        name="scores", type="array", description="List of scores", items=item_schema
+    )
+
+    param = schema.to_param()
+    assert isinstance(param, Parameter)
+    assert param.name == "scores"
+    assert param.annotation == list[int]
+    assert param.kind == Parameter.POSITIONAL_OR_KEYWORD
+
+
+def test_parameter_schema_array_no_items_error():
+    """Tests that 'array' type raises error if 'items' is None."""
+    schema = ParameterSchema(
+        name="bad_list", type="array", description="List without item type"
+    )
+
+    expected_error_msg = "Unexpected value: type is 'list' but items is None"
+    with pytest.raises(Exception, match=expected_error_msg):
+        schema._ParameterSchema__get_type()
+
+    with pytest.raises(Exception, match=expected_error_msg):
+        schema.to_param()
+
+
+def test_parameter_schema_unsupported_type_error():
+    """Tests that an unsupported type raises ValueError."""
+    unsupported_type = "datetime"
+    schema = ParameterSchema(
+        name="event_time", type=unsupported_type, description="When it happened"
+    )
+
+    expected_error_msg = f"Unsupported schema type: {unsupported_type}"
+    with pytest.raises(ValueError, match=expected_error_msg):
+        schema._ParameterSchema__get_type()
+
+    with pytest.raises(ValueError, match=expected_error_msg):
+        schema.to_param()