fix(tests): resolve unit test failures in tool and client wrappers

Author: anubhav756

packages/toolbox-adk/src/toolbox_adk/tool.py

@@ -12,25 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Optional, Callable, Awaitable
-from typing_extensions import override
+from typing import Any, Awaitable, Callable, Dict, Optional
 
 import toolbox_core
-from toolbox_core.tool import ToolboxTool as CoreToolboxTool
-from google.adk.tools.base_tool import BaseTool
+from fastapi.openapi.models import (
+    OAuth2,
+    OAuthFlowAuthorizationCode,
+    OAuthFlows,
+    SecurityScheme,
+)
 from google.adk.agents.readonly_context import ReadonlyContext
-
+from google.adk.auth.auth_credential import (
+    AuthCredential,
+    AuthCredentialTypes,
+    OAuth2Auth,
+)
 from google.adk.auth.auth_tool import AuthConfig
-from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes, OAuth2Auth
-from fastapi.openapi.models import SecurityScheme, OAuthFlows, OAuthFlowAuthorizationCode, OAuth2
+from google.adk.tools.base_tool import BaseTool
+from toolbox_core.tool import ToolboxTool as CoreToolboxTool
+from typing_extensions import override
 
-from .credentials import CredentialConfig, CredentialType
 from .client import USER_TOKEN_CONTEXT_VAR
-
+from .credentials import CredentialConfig, CredentialType
 
 
 class ToolboxContext:
     """Context object passed to pre/post hooks."""
+
     def __init__(self, arguments: Dict[str, Any], tool_context: ReadonlyContext):
         self.arguments = arguments
         self.tool_context = tool_context
@@ -59,15 +67,18 @@ def __init__(
         """
         # We act as a proxy.
         # We need to extract metadata from the core tool to satisfy BaseTool's contract.
-        
+
         name = getattr(core_tool, "__name__", "unknown_tool")
-        description = getattr(core_tool, "__doc__", "No description provided.") or "No description provided."
-        
+        description = (
+            getattr(core_tool, "__doc__", "No description provided.")
+            or "No description provided."
+        )
+
         super().__init__(
             name=name,
             description=description,
             # We pass empty custom_metadata or whatever is needed
-            custom_metadata={} 
+            custom_metadata={},
         )
         self._core_tool = core_tool
         self._pre_hook = pre_hook
@@ -82,30 +93,32 @@ async def run_async(
     ) -> Any:
         # Create context
         ctx = ToolboxContext(arguments=args, tool_context=tool_context)
-        
+
         # 1. Pre-hook
         if self._pre_hook:
             await self._pre_hook(ctx)
-            
+
         # 2. ADK Auth Integration (3LO)
         # Check if USER_IDENTITY is configured
         reset_token = None
-        
+
         if self._auth_config and self._auth_config.type == CredentialType.USER_IDENTITY:
             if not self._auth_config.client_id or not self._auth_config.client_secret:
                 raise ValueError("USER_IDENTITY requires client_id and client_secret")
-                
+
             # Construct ADK AuthConfig
-            scopes = self._auth_config.scopes or ["https://www.googleapis.com/auth/cloud-platform"]
+            scopes = self._auth_config.scopes or [
+                "https://www.googleapis.com/auth/cloud-platform"
+            ]
             scope_dict = {s: "" for s in scopes}
-            
+
             auth_config_adk = AuthConfig(
                 auth_scheme=OAuth2(
                     flows=OAuthFlows(
                         authorizationCode=OAuthFlowAuthorizationCode(
                             authorizationUrl="https://accounts.google.com/o/oauth2/auth",
                             tokenUrl="https://oauth2.googleapis.com/token",
-                            scopes=scope_dict
+                            scopes=scope_dict,
                         )
                     )
                 ),
@@ -114,9 +127,9 @@ async def run_async(
                     oauth2=OAuth2Auth(
                         client_id=self._auth_config.client_id,
                         client_secret=self._auth_config.client_secret,
-                        scopes=scopes
-                    )
-                )
+                        scopes=scopes,
+                    ),
+                ),
             )
 
             # Check if we already have credentials from a previous exchange
@@ -140,19 +153,21 @@ async def run_async(
                 # Actually, strictly we should probably request credential if get_auth_response returns nothing
                 # but get_auth_response typically handles the lookup.
                 # If exception is unrelated, raise.
-                if "credential" in str(e).lower() or isinstance(e, ValueError): # Loose check, but safest is to re-raise
-                     raise e
+                if "credential" in str(e).lower() or isinstance(
+                    e, ValueError
+                ):  # Loose check, but safest is to re-raise
+                    raise e
                 # Fallback to request logic if it was a lookup error?
                 tool_context.request_credential(auth_config_adk)
                 return None
 
         try:
             # Execute the core tool
             result = await self._core_tool(**ctx.arguments)
-            
+
             ctx.result = result
             return result
-            
+
         except Exception as e:
             ctx.error = e
             raise e
@@ -162,13 +177,13 @@ async def run_async(
             if self._post_hook:
                 await self._post_hook(ctx)
 
-    def bind_params(self, bounded_params: Dict[str, Any]) -> 'ToolboxTool':
+    def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool":
         """Allows runtime binding of parameters, delegating to core tool."""
         new_core_tool = self._core_tool.bind_params(bounded_params)
         # Return a new wrapper
         return ToolboxTool(
-            core_tool=new_core_tool, 
-            pre_hook=self._pre_hook, 
+            core_tool=new_core_tool,
+            pre_hook=self._pre_hook,
             post_hook=self._post_hook,
-            auth_config=self._auth_config
+            auth_config=self._auth_config,
         )

packages/toolbox-adk/src/toolbox_adk/toolset.py

@@ -12,16 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Optional, Mapping, Union, Callable, Awaitable
-from typing_extensions import override
+from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Union
 
+from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.tools.base_tool import BaseTool
 from google.adk.tools.base_toolset import BaseToolset
-from google.adk.agents.readonly_context import ReadonlyContext
+from typing_extensions import override
 
 from .client import ToolboxClient
 from .credentials import CredentialConfig
-from .tool import ToolboxTool, ToolboxContext
+from .tool import ToolboxContext, ToolboxTool
+
 
 class ToolboxToolset(BaseToolset):
     """
@@ -34,11 +35,13 @@ def __init__(
         toolset_name: Optional[str] = None,
         tool_names: Optional[List[str]] = None,
         credentials: Optional[CredentialConfig] = None,
-        additional_headers: Optional[Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]]] = None,
+        additional_headers: Optional[
+            Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]]
+        ] = None,
         bound_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None,
         pre_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
         post_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
-        **kwargs
+        **kwargs,
     ):
         """
         Args:
@@ -56,7 +59,7 @@ def __init__(
             server_url=server_url,
             credentials=credentials,
             additional_headers=additional_headers,
-            **kwargs
+            **kwargs,
         )
         self._toolset_name = toolset_name
         self._tool_names = tool_names
@@ -69,38 +72,36 @@ async def get_tools(
         self, readonly_context: Optional[ReadonlyContext] = None
     ) -> List[BaseTool]:
         """Loads tools from the toolbox server and wraps them."""
-        # Note: We don't close the client after get_tools because tools might need it 
+        # Note: We don't close the client after get_tools because tools might need it
         # (though currently tools are self-contained http-wise if they don't share session state).
         # Actually toolbox-core tools have their own client reference or use a shared one?
-        # toolbox_core.ToolboxClient.load_tool returns a tool which... 
+        # toolbox_core.ToolboxClient.load_tool returns a tool which...
         # wait, toolbox_core tools make HTTP calls. Do they hold a reference to the client session?
         # Yes, toolbox_core 0.1.0+ tools usually have an internal `client` or `session`.
         # So we must keep self._client alive or ensuring it's not closed prematurely.
-        
+
         tools = []
         if self._toolset_name:
             core_tools = await self._client.load_toolset(
-                self._toolset_name,
-                bound_params=self._bound_params
+                self._toolset_name, bound_params=self._bound_params
             )
             tools.extend(core_tools)
 
         if self._tool_names:
             for name in self._tool_names:
                 core_tool = await self._client.load_tool(
-                    name,
-                    bound_params=self._bound_params
+                    name, bound_params=self._bound_params
                 )
                 tools.append(core_tool)
-        
+
         # Wrap all core tools in ToolboxTool
         return [
             ToolboxTool(
-                core_tool=t, 
-                pre_hook=self._pre_hook, 
+                core_tool=t,
+                pre_hook=self._pre_hook,
                 post_hook=self._post_hook,
-                auth_config=self._client.credential_config
-            ) 
+                auth_config=self._client.credential_config,
+            )
             for t in tools
         ]
 

packages/toolbox-adk/tests/unit/test_client.py

@@ -27,17 +27,15 @@ def test_init_no_auth(self, mock_core_client):
         creds = CredentialStrategy.TOOLBOX_IDENTITY()
         client = ToolboxClient("http://server", credentials=creds)
 
-        mock_core_client.assert_called_with(
-            server_url="http://server", client_headers={}
-        )
+        mock_core_client.assert_called_with(url="http://server", client_headers={})
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
     def test_init_manual_token(self, mock_core_client):
         creds = CredentialStrategy.MANUAL_TOKEN("abc")
         client = ToolboxClient("http://server", credentials=creds)
 
         mock_core_client.assert_called_with(
-            server_url="http://server", client_headers={"Authorization": "Bearer abc"}
+            url="http://server", client_headers={"Authorization": "Bearer abc"}
         )
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
@@ -49,7 +47,7 @@ def test_init_additional_headers(self, mock_core_client):
         )
 
         mock_core_client.assert_called_with(
-            server_url="http://server", client_headers={"X-Custom": "Val"}
+            url="http://server", client_headers={"X-Custom": "Val"}
         )
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")