test(client): add comprehensive unit tests for 100% coverage

anubhav756 · anubhav756 · commit 779d9cd241b1 · 2025-12-11T03:51:00.000+05:30
diff --git a/packages/toolbox-adk/src/toolbox_adk/client.py b/packages/toolbox-adk/src/toolbox_adk/client.py
@@ -22,7 +22,6 @@
 
 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
 )
@@ -53,11 +52,6 @@ def __init__(
         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], Callable[[], Awaitable[str]]]
         ] = {}
@@ -85,8 +79,6 @@ def _configure_auth(self, creds: CredentialConfig) -> None:
                 )
 
             # 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
             )
@@ -108,14 +100,10 @@ def _configure_auth(self, creds: CredentialConfig) -> None:
 
         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
-                    # we expected it. We return empty string which might cause 401.
                     return ""
                 return f"Bearer {token}"
 
@@ -125,28 +113,19 @@ 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.
             request = transport.requests.Request()
-            # Try to get ID token directly (e.g. on Cloud Run)
             try:
                 token = id_token.fetch_id_token(request, audience)
                 return f"Bearer {token}"
             except Exception:
-                # Fallback to default credentials (e.g. local gcloud)
+                # Fallback to default credentials
                 creds, _ = google.auth.default()
                 if not creds.valid:
                     creds.refresh(request)
-                # If specific ID token credentials, use them, otherwise this might be Access Token (scoped)
-                # 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 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
                 curr_token = getattr(creds, "token", None)
                 return f"Bearer {curr_token}" if curr_token else ""
 
diff --git a/packages/toolbox-adk/tests/unit/test_client.py b/packages/toolbox-adk/tests/unit/test_client.py
@@ -12,159 +12,178 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import ANY, AsyncMock, MagicMock, patch
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from toolbox_adk.client import ToolboxClient
-from toolbox_adk.credentials import CredentialStrategy
+from toolbox_adk import CredentialStrategy, ToolboxClient
+from toolbox_adk.client import CredentialType
 
 
-class TestToolboxClient:
+@pytest.mark.asyncio
+class TestToolboxClientAuth:
+    """Unit tests for Client Auth logic."""
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
-    def test_init_no_auth(self, mock_core_client):
+    async def test_init_toolbox_identity(self, mock_core_client):
+        """Test init with TOOLBOX_IDENTITY (no auth headers)."""
         creds = CredentialStrategy.TOOLBOX_IDENTITY()
-        client = ToolboxClient("http://server", credentials=creds)
+        client = ToolboxClient(server_url="http://test", credentials=creds)
 
-        mock_core_client.assert_called_with(
-            server_url="http://server", client_headers={}
-        )
+        # Verify core client created with empty headers for auth
+        _, kwargs = mock_core_client.call_args
+        assert "client_headers" in kwargs
+        headers = kwargs["client_headers"]
+        assert "Authorization" not in 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)
+    @patch("toolbox_adk.client.id_token.fetch_id_token")
+    @patch("toolbox_adk.client.google.auth.default")
+    @patch("toolbox_adk.client.transport.requests.Request")
+    async def test_init_adc_success_fetch_id_token(
+        self, mock_req, mock_default, mock_fetch_id, mock_core_client
+    ):
+        """Test ADC strategy where fetch_id_token succeeds."""
+        mock_fetch_id.return_value = "id-token-123"
 
-        mock_core_client.assert_called_with(
-            server_url="http://server", client_headers={"Authorization": "Bearer abc"}
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(
+            target_audience="aud"
         )
+        client = ToolboxClient(server_url="http://test", credentials=creds)
 
-    @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
-        )
+        _, kwargs = mock_core_client.call_args
+        headers = kwargs["client_headers"]
+        assert "Authorization" in headers
+        token_getter = headers["Authorization"]
+        assert callable(token_getter)
 
-        mock_core_client.assert_called_with(
-            server_url="http://server", client_headers={"X-Custom": "Val"}
-        )
+        # Call the getter
+        token_val = token_getter()
+        assert token_val == "Bearer id-token-123"
+        mock_fetch_id.assert_called()
 
     @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"
+    @patch("toolbox_adk.client.google.auth.default")
+    @patch("toolbox_adk.client.transport.requests.Request")
+    async def test_init_adc_fallback_creds(
+        self, mock_req, mock_default, mock_fetch_id, mock_core_client
+    ):
+        """Test ADC strategy fallback to default() when fetch_id_token fails."""
+        mock_fetch_id.side_effect = Exception("No metadata server")
 
-        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
-        client = ToolboxClient("http://server", credentials=creds)
+        # Mock default creds
+        mock_creds = MagicMock()
+        mock_creds.valid = False
+        mock_creds.id_token = "fallback-id-token"
+        mock_default.return_value = (mock_creds, "proj")
 
-        # 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)
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(
+            target_audience="aud"
+        )
+        client = ToolboxClient(server_url="http://test", credentials=creds)
 
-        # Verify callable behavior
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
         token = token_getter()
-        assert token == "Bearer id_token_123"
-        mock_fetch_token.assert_called()
+        assert token == "Bearer fallback-id-token"
+        mock_creds.refresh.assert_called()  # Because we set valid=False
 
     @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
+    @patch("toolbox_adk.client.transport.requests.Request")
+    async def test_init_adc_fallback_creds_token(
+        self, mock_req, mock_default, mock_fetch_id, mock_core_client
     ):
-        # unexpected error on fetch_id_token
-        mock_fetch_token.side_effect = Exception("No metadata")
+        """Test ADC fallback when creds have .token but no .id_token."""
+        mock_fetch_id.side_effect = Exception("No metadata server")
 
         mock_creds = MagicMock()
-        mock_creds.valid = False
-        mock_creds.id_token = "fallback_id_token"
+        mock_creds.valid = True
+        del mock_creds.id_token  # Simulate no id_token attr or None
+        mock_creds.token = "access-token-123"  # e.g. user creds
         mock_default.return_value = (mock_creds, "proj")
 
-        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
-        client = ToolboxClient("http://server", credentials=creds)
-
+        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(
+            target_audience="aud"
+        )
+        client = ToolboxClient(server_url="http://test", credentials=creds)
         token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
-        token = token_getter()
+        assert token_getter() == "Bearer access-token-123"
 
-        assert token == "Bearer fallback_id_token"
-        mock_creds.refresh.assert_called()
+    @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
+    async def test_init_manual_token(self, mock_core_client):
+        creds = CredentialStrategy.MANUAL_TOKEN(token="abc")
+        client = ToolboxClient("http://test", credentials=creds)
+        headers = mock_core_client.call_args[1]["client_headers"]
+        assert headers["Authorization"] == "Bearer abc"
 
     @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"
+    async def test_init_manual_creds(self, mock_core_client):
+        mock_google_creds = MagicMock()
+        mock_google_creds.valid = True
+        mock_google_creds.token = "creds-token"
 
-        creds = CredentialStrategy.MANUAL_CREDS(mock_g_creds)
-        client = ToolboxClient("http://server", credentials=creds)
+        creds = CredentialStrategy.MANUAL_CREDS(credentials=mock_google_creds)
+        client = ToolboxClient("http://test", 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()
+        assert token_getter() == "Bearer creds-token"
 
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
-    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),
-            )
-
-        # 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))
+    async def test_init_user_identity(self, mock_core_client):
+        creds = CredentialStrategy.USER_IDENTITY(client_id="c", client_secret="s")
+        client = ToolboxClient("http://test", credentials=creds)
 
-    @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
-    ):
-        # fetch_id_token fails
-        mock_fetch_token.side_effect = Exception("No metadata")
+        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
+        # Should be empty initially
+        assert token_getter() == ""
 
-        mock_creds = MagicMock()
-        mock_creds.valid = False
-        mock_creds.id_token = None  # No ID token
-        mock_creds.token = "access_token_123"
-        mock_default.return_value = (mock_creds, "proj")
+        # Set context
+        from toolbox_adk.client import USER_TOKEN_CONTEXT_VAR
 
-        creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS("http://aud")
-        client = ToolboxClient("http://server", credentials=creds)
+        token = USER_TOKEN_CONTEXT_VAR.set("user-tok")
+        try:
+            assert token_getter() == "Bearer user-tok"
+        finally:
+            USER_TOKEN_CONTEXT_VAR.reset(token)
 
-        token_getter = mock_core_client.call_args[1]["client_headers"]["Authorization"]
-        token = token_getter()
+    async def test_validation_errors(self):
+        with pytest.raises(ValueError):
+            # ADC requires audience
+            creds = CredentialStrategy.APPLICATION_DEFAULT_CREDENTIALS(target_audience="")
+            ToolboxClient("http://test", credentials=creds)
 
-        assert token == "Bearer access_token_123"
-        mock_creds.refresh.assert_called()
+        with pytest.raises(ValueError):
+            creds = CredentialStrategy.MANUAL_TOKEN(token="")
+            ToolboxClient("http://test", credentials=creds)
+
+        with pytest.raises(ValueError):
+            creds = CredentialStrategy.MANUAL_CREDS(credentials=None)
+            ToolboxClient("http://test", credentials=creds)
 
-    @pytest.mark.asyncio
     @patch("toolbox_adk.client.toolbox_core.ToolboxClient")
-    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()
+    async def test_load_methods(self, mock_core_client_class):
+        # Setup mock instance
+        mock_instance = AsyncMock()
+        mock_core_client_class.return_value = mock_instance
 
-        client = ToolboxClient("http://server")
-        tools = await client.load_toolset("my-set", extra="arg")
+        client = ToolboxClient(
+            "http://test", credentials=CredentialStrategy.TOOLBOX_IDENTITY()
+        )
 
-        mock_instance.load_toolset.assert_awaited_with("my-set", extra="arg")
-        assert tools == ["t1"]
+        # Test load_toolset
+        await client.load_toolset("ts", foo="bar")
+        mock_instance.load_toolset.assert_called_with("ts", foo="bar")
 
+        # Test load_tool
+        await client.load_tool("t", baz="qux")
+        mock_instance.load_tool.assert_called_with("t", baz="qux")
+
+        # Test close
         await client.close()
-        mock_instance.close.assert_awaited()
+        mock_instance.close.assert_called_once()
+
+        # Test property
+        assert client.credential_config is not None
