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

Author: anubhav756

.github/workflows/lint-toolbox-adk.yaml

@@ -0,0 +1,92 @@
+# 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.
+
+name: adk
+on:
+  pull_request:
+    paths:
+      - 'packages/toolbox-adk/**'
+      - '!packages/toolbox-adk/**/*.md'
+  pull_request_target:
+    types: [labeled]
+
+# Declare default permissions as read only.
+permissions: read-all
+
+jobs:
+  lint:
+    if: "${{ github.event.action != 'labeled' || github.event.label.name == 'tests: run' }}"
+    name: lint
+    runs-on: ubuntu-latest
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.ref }}
+      cancel-in-progress: true
+    defaults:
+      run:
+        working-directory: ./packages/toolbox-adk
+    permissions:
+      contents: 'read'
+      issues: 'write'
+      pull-requests: 'write'
+    steps:
+      - name: Remove PR Label
+        if: "${{ github.event.action == 'labeled' && github.event.label.name == 'tests: run' }}"
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            try {
+              await github.rest.issues.removeLabel({
+                name: 'tests: run',
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                issue_number: context.payload.pull_request.number
+              });
+            } catch (e) {
+              console.log('Failed to remove label. Another job may have already removed it!');
+            }
+      - name: Checkout code
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+          repository: ${{ github.event.pull_request.head.repo.full_name }}
+          token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Setup Python
+        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+        with:
+          python-version: "3.11"
+
+      - name: Install dependencies
+        run: |
+          pip install .[test]
+          # Install toolbox-core from local path to ensure we test against current codebase
+          pip install ../toolbox-core
+
+      - name: Run linters
+        run: |
+          # Check if we need to run linters or just rely on pre-commit locally?
+          # The other files allow this.
+          # For consistency:
+          # black --check .
+          # isort --check .
+          # But let's check if the user wants strictly same contents or just 'tests'.
+          # I will add basic linting to be safe/compliant.
+          pip install black isort mypy
+          black --check .
+          isort --check .
+
+      - name: Run type-check
+        env:
+          MYPYPATH: './src'
+        run: mypy --install-types --non-interactive --cache-dir=.mypy_cache/ -p toolbox_adk

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 ""