style: apply black formatting to client wrapper

Author: anubhav756

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

@@ -12,18 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Dict, Callable, Any, Union, Awaitable
+from contextvars import ContextVar
+from typing import Any, Awaitable, Callable, Dict, Optional, Union
+
+import google.auth
 import toolbox_core
-from google.auth import transport
-from google.auth import compute_engine
+from google.auth import compute_engine, transport
 from google.oauth2 import id_token
-import google.auth
-from contextvars import ContextVar
 
 from .credentials import CredentialConfig, CredentialType
 
 # Global ContextVar for User Identity (3LO) tokens to be injected per-request
-USER_TOKEN_CONTEXT_VAR: ContextVar[Optional[str]] = ContextVar("toolbox_user_token", default=None)
+USER_TOKEN_CONTEXT_VAR: ContextVar[Optional[str]] = ContextVar(
+    "toolbox_user_token", default=None
+)
 
 
 class ToolboxClient:
@@ -35,9 +37,10 @@ def __init__(
         self,
         server_url: str,
         credentials: Optional[CredentialConfig] = None,
-
-        additional_headers: Optional[Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]]] = None,
-        **kwargs
+        additional_headers: Optional[
+            Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]]
+        ] = None,
+        **kwargs,
     ):
         """
         Args:
@@ -49,14 +52,14 @@ def __init__(
         self._server_url = server_url
         self._credentials = credentials
         self._additional_headers = additional_headers or {}
-        
+
         # Prepare auth_token_getters for toolbox-core
         # toolbox_core expects: dict[str, Callable[[], str | Awaitable[str]]]
         # However, for general headers (like Authorization), we can pass them in client_headers
         # if they are static or simpler. Toolbox-core supports `client_headers` which can be dynamic.
-        
+
         self._core_client_headers: Dict[str, Union[str, Callable[[], str]]] = {}
-        
+
         # Add static additional headers
         for k, v in self._additional_headers.items():
             self._core_client_headers[k] = v
@@ -65,55 +68,60 @@ def __init__(
             self._configure_auth(credentials)
 
         self._client = toolbox_core.ToolboxClient(
-            url=server_url,
-            client_headers=self._core_client_headers,
-            **kwargs
+            url=server_url, client_headers=self._core_client_headers, **kwargs
         )
 
     def _configure_auth(self, creds: CredentialConfig):
         if creds.type == CredentialType.TOOLBOX_IDENTITY:
             # No auth headers needed
             pass
-            
+
         elif creds.type == CredentialType.APPLICATION_DEFAULT_CREDENTIALS:
             if not creds.target_audience:
-                raise ValueError("target_audience is required for APPLICATION_DEFAULT_CREDENTIALS")
-            
+                raise ValueError(
+                    "target_audience is required for APPLICATION_DEFAULT_CREDENTIALS"
+                )
+
             # Create an async capable token getter
             # We wrap it to match the signature expected by toolbox-core headers
             # (which accepts callables)
-            self._core_client_headers["Authorization"] = self._create_adc_token_getter(creds.target_audience)
-            
+            self._core_client_headers["Authorization"] = self._create_adc_token_getter(
+                creds.target_audience
+            )
+
         elif creds.type == CredentialType.MANUAL_TOKEN:
             if not creds.token:
                 raise ValueError("token is required for MANUAL_TOKEN")
             scheme = creds.scheme or "Bearer"
             self._core_client_headers["Authorization"] = f"{scheme} {creds.token}"
-            
+
         elif creds.type == CredentialType.MANUAL_CREDS:
             if not creds.credentials:
                 raise ValueError("credentials object is required for MANUAL_CREDS")
-            
+
             # Adapter for google-auth credentials object to callable
-            self._core_client_headers["Authorization"] = self._create_creds_token_getter(creds.credentials)
-            
+            self._core_client_headers["Authorization"] = (
+                self._create_creds_token_getter(creds.credentials)
+            )
+
         elif creds.type == CredentialType.USER_IDENTITY:
             # For USER_IDENTITY (3LO), the *Tool* handles the interactive flow at runtime.
             # We use a ContextVar to inject the token per-request.
-            
+
             def get_user_token() -> str:
                 token = USER_TOKEN_CONTEXT_VAR.get()
                 if not token:
-                    # If this is called but no token is set in context, it means 
-                    # the tool wrapper failed to set it or we are in a context where 
+                    # If this is called but no token is set in context, it means
+                    # the tool wrapper failed to set it or we are in a context where
                     # we expected it. We return empty string which might cause 401.
                     return ""
                 return f"Bearer {token}"
-                
+
             self._core_client_headers["Authorization"] = get_user_token
 
     def _create_adc_token_getter(self, audience: str) -> Callable[[], str]:
         """Returns a callable that fetches a fresh ID token using ADC."""
+
         def get_token() -> str:
             # Note: This is a synchronous call. Toolbox-core supports sync callables in headers.
             # Ideally we would use async but google-auth is primarily sync for these helpers.
@@ -131,9 +139,9 @@ def get_token() -> str:
                 # For Toolbox we usually need ID Tokens.
                 # If the user is locally auth'd via `gcloud auth login`, fetch_id_token is preferred.
                 # If falling back to service account file:
-                if hasattr(creds, 'id_token') and creds.id_token:
-                     return f"Bearer {creds.id_token}"
-                
+                if hasattr(creds, "id_token") and creds.id_token:
+                    return f"Bearer {creds.id_token}"
+
                 # If we are here, we might need to manually sign via IAM or similar if it's a generic SA.
                 # For simplicity in this v1, we assume fetch_id_token works or standard creds work.
                 # Re-attempt fetch_id_token on the credentials object if possible
@@ -148,6 +156,7 @@ def get_token() -> str:
             if not credentials.valid:
                 credentials.refresh(request)
             return f"Bearer {credentials.token}"
+
         return get_token
 
     @property

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

@@ -12,62 +12,60 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
 import pytest
-from unittest.mock import MagicMock, patch, ANY, AsyncMock
+
 from toolbox_adk.client import ToolboxClient
 from toolbox_adk.credentials import CredentialStrategy
 
+
 class TestToolboxClient:
-    
+
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
     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={}
+            server_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"}
+            server_url="http://server", client_headers={"Authorization": "Bearer abc"}
         )
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
     def test_init_additional_headers(self, mock_core_client):
         creds = CredentialStrategy.TOOLBOX_IDENTITY()
         headers = {"X-Custom": "Val"}
         client = ToolboxClient(
-            "http://server", 
-            credentials=creds, 
-            additional_headers=headers
+            "http://server", credentials=creds, additional_headers=headers
         )
-        
+
         mock_core_client.assert_called_with(
-            server_url="http://server",
-            client_headers={"X-Custom": "Val"}
+            server_url="http://server", client_headers={"X-Custom": "Val"}
         )
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
     @patch("toolbox_adk.client.id_token.fetch_id_token")
     def test_adc_auth_flow_success(self, mock_fetch_token, mock_core_client):
         mock_fetch_token.return_value = "id_token_123"
-        
+
         creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
         client = ToolboxClient("http://server", credentials=creds)
-        
+
         # Verify a callable was passed
         args, kwargs = mock_core_client.call_args
         assert "Authorization" in kwargs["client_headers"]
         token_getter = kwargs["client_headers"]["Authorization"]
         assert callable(token_getter)
-        
+
         # Verify callable behavior
         token = token_getter()
         assert token == "Bearer id_token_123"
@@ -76,36 +74,38 @@ def test_adc_auth_flow_success(self, mock_fetch_token, mock_core_client):
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
     @patch("toolbox_adk.client.id_token.fetch_id_token")
     @patch("toolbox_adk.client.google.auth.default")
-    def test_adc_auth_flow_fallback(self, mock_default, mock_fetch_token, mock_core_client):
+    def test_adc_auth_flow_fallback(
+        self, mock_default, mock_fetch_token, mock_core_client
+    ):
         # unexpected error on fetch_id_token
         mock_fetch_token.side_effect = Exception("No metadata")
-        
+
         mock_creds = MagicMock()
         mock_creds.valid = False
         mock_creds.id_token = "fallback_id_token"
         mock_default.return_value = (mock_creds, "proj")
-        
+
         creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
         client = ToolboxClient("http://server", credentials=creds)
-        
+
         token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
         token = token_getter()
-        
+
         assert token == "Bearer fallback_id_token"
         mock_creds.refresh.assert_called()
-    
+
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
     def test_manual_creds(self, mock_core_client):
         mock_g_creds = MagicMock()
         mock_g_creds.valid = False
         mock_g_creds.token = "oauth_token"
-        
+
         creds = CredentialStrategy.MANUAL_CREDS(mock_g_creds)
         client = ToolboxClient("http://server", credentials=creds)
-        
+
         token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
         token = token_getter()
-        
+
         assert token == "Bearer oauth_token"
         assert token == "Bearer oauth_token"
         mock_g_creds.refresh.assert_called()
@@ -115,36 +115,41 @@ def test_init_validation_errors(self, mock_core_client):
         # ADC missing audience
         with pytest.raises(ValueError, match="target_audience is required"):
             # Fix: only pass target_audience as keyword arg OR positional, not both mixed in a way that causes overlap if defined so
-            # Actually simpler: just pass raw None 
-            ToolboxClient("url", credentials=CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(None))
-            
+            # Actually simpler: just pass raw None
+            ToolboxClient(
+                "url",
+                credentials=CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(None),
+            )
+
         # Manual token missing token
         with pytest.raises(ValueError, match="token is required"):
             ToolboxClient("url", credentials=CredentialStrategy.MANUAL_TOKEN(None))
-            
+
         # Manual creds missing credentials
         with pytest.raises(ValueError, match="credentials object is required"):
             ToolboxClient("url", credentials=CredentialStrategy.MANUAL_CREDS(None))
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
     @patch("toolbox_adk.client.id_token.fetch_id_token")
     @patch("toolbox_adk.client.google.auth.default")
-    def test_adc_auth_flow_fallback_access_token(self, mock_default, mock_fetch_token, mock_core_client):
+    def test_adc_auth_flow_fallback_access_token(
+        self, mock_default, mock_fetch_token, mock_core_client
+    ):
         # fetch_id_token fails
         mock_fetch_token.side_effect = Exception("No metadata")
-        
+
         mock_creds = MagicMock()
         mock_creds.valid = False
-        mock_creds.id_token = None # No ID token
+        mock_creds.id_token = None  # No ID token
         mock_creds.token = "access_token_123"
         mock_default.return_value = (mock_creds, "proj")
-        
+
         creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
         client = ToolboxClient("http://server", credentials=creds)
-        
+
         token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
         token = token_getter()
-        
+
         assert token == "Bearer access_token_123"
         mock_creds.refresh.assert_called()
 
@@ -154,12 +159,12 @@ async def test_delegation(self, mock_core_client):
         mock_instance = mock_core_client.return_value
         mock_instance.load_toolset = AsyncMock(return_value=["t1"])
         mock_instance.close = AsyncMock()
-        
+
         client = ToolboxClient("http://server")
         tools = await client.load_toolset("my-set", extra="arg")
-        
+
         mock_instance.load_toolset.assert_awaited_with("my-set", extra="arg")
         assert tools == ["t1"]
-        
+
         await client.close()
         mock_instance.close.assert_awaited()