Moves unittests to root folder and adds github action to run unit tests. (#72)

* Move unit tests to root package.

* Adds deps to "test" extra, and mark two broken tests in tests/unittests/auth/test_auth_handler.py

* Adds github workflow

* minor fix in lite_llm.py for python 3.9.

* format pyproject.toml
This commit is contained in:
Jack Sun
2025-04-11 08:25:59 -07:00
committed by GitHub
parent 59117b9b96
commit 05142a07cc
66 changed files with 50 additions and 2 deletions

View File

@@ -0,0 +1,145 @@
# 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.
"""Unit tests for AutoAuthCredentialExchanger."""
from typing import Dict
from typing import Optional
from typing import Type
from unittest.mock import MagicMock
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_schemes import AuthScheme
from google.adk.tools.openapi_tool.auth.credential_exchangers.auto_auth_credential_exchanger import AutoAuthCredentialExchanger
from google.adk.tools.openapi_tool.auth.credential_exchangers.base_credential_exchanger import BaseAuthCredentialExchanger
from google.adk.tools.openapi_tool.auth.credential_exchangers.oauth2_exchanger import OAuth2CredentialExchanger
from google.adk.tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger
import pytest
class MockCredentialExchanger(BaseAuthCredentialExchanger):
"""Mock credential exchanger for testing."""
def exchange_credential(
self,
auth_scheme: AuthScheme,
auth_credential: Optional[AuthCredential] = None,
) -> AuthCredential:
"""Mock exchange credential method."""
return auth_credential
@pytest.fixture
def auto_exchanger():
"""Fixture for creating an AutoAuthCredentialExchanger instance."""
return AutoAuthCredentialExchanger()
@pytest.fixture
def auth_scheme():
"""Fixture for creating a mock AuthScheme instance."""
scheme = MagicMock(spec=AuthScheme)
return scheme
def test_init_with_custom_exchangers():
"""Test initialization with custom exchangers."""
custom_exchangers: Dict[str, Type[BaseAuthCredentialExchanger]] = {
AuthCredentialTypes.API_KEY: MockCredentialExchanger
}
auto_exchanger = AutoAuthCredentialExchanger(
custom_exchangers=custom_exchangers
)
assert (
auto_exchanger.exchangers[AuthCredentialTypes.API_KEY]
== MockCredentialExchanger
)
assert (
auto_exchanger.exchangers[AuthCredentialTypes.OPEN_ID_CONNECT]
== OAuth2CredentialExchanger
)
def test_exchange_credential_no_auth_credential(auto_exchanger, auth_scheme):
"""Test exchange_credential with no auth_credential."""
assert auto_exchanger.exchange_credential(auth_scheme, None) is None
def test_exchange_credential_no_exchange(auto_exchanger, auth_scheme):
"""Test exchange_credential with NoExchangeCredentialExchanger."""
auth_credential = AuthCredential(auth_type=AuthCredentialTypes.API_KEY)
result = auto_exchanger.exchange_credential(auth_scheme, auth_credential)
assert result == auth_credential
def test_exchange_credential_open_id_connect(auto_exchanger, auth_scheme):
"""Test exchange_credential with OpenID Connect scheme."""
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.OPEN_ID_CONNECT
)
mock_exchanger = MagicMock(spec=OAuth2CredentialExchanger)
mock_exchanger.exchange_credential.return_value = "exchanged_credential"
auto_exchanger.exchangers[AuthCredentialTypes.OPEN_ID_CONNECT] = (
lambda: mock_exchanger
)
result = auto_exchanger.exchange_credential(auth_scheme, auth_credential)
assert result == "exchanged_credential"
mock_exchanger.exchange_credential.assert_called_once_with(
auth_scheme, auth_credential
)
def test_exchange_credential_service_account(auto_exchanger, auth_scheme):
"""Test exchange_credential with Service Account scheme."""
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT
)
mock_exchanger = MagicMock(spec=ServiceAccountCredentialExchanger)
mock_exchanger.exchange_credential.return_value = "exchanged_credential_sa"
auto_exchanger.exchangers[AuthCredentialTypes.SERVICE_ACCOUNT] = (
lambda: mock_exchanger
)
result = auto_exchanger.exchange_credential(auth_scheme, auth_credential)
assert result == "exchanged_credential_sa"
mock_exchanger.exchange_credential.assert_called_once_with(
auth_scheme, auth_credential
)
def test_exchange_credential_custom_exchanger(auto_exchanger, auth_scheme):
"""Test that exchange_credential calls the correct (custom) exchanger."""
# Use a custom exchanger via the initialization
mock_exchanger = MagicMock(spec=MockCredentialExchanger)
mock_exchanger.exchange_credential.return_value = "custom_credential"
auto_exchanger.exchangers[AuthCredentialTypes.API_KEY] = (
lambda: mock_exchanger
)
auth_credential = AuthCredential(auth_type=AuthCredentialTypes.API_KEY)
result = auto_exchanger.exchange_credential(auth_scheme, auth_credential)
assert result == "custom_credential"
mock_exchanger.exchange_credential.assert_called_once_with(
auth_scheme, auth_credential
)

View File

@@ -0,0 +1,68 @@
# 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.
"""Tests for the BaseAuthCredentialExchanger class."""
from typing import Optional
from unittest.mock import MagicMock
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_schemes import AuthScheme
from google.adk.tools.openapi_tool.auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError
from google.adk.tools.openapi_tool.auth.credential_exchangers.base_credential_exchanger import BaseAuthCredentialExchanger
import pytest
class MockAuthCredentialExchanger(BaseAuthCredentialExchanger):
def exchange_credential(
self,
auth_scheme: AuthScheme,
auth_credential: Optional[AuthCredential] = None,
) -> AuthCredential:
return AuthCredential(token="some-token")
class TestBaseAuthCredentialExchanger:
"""Tests for the BaseAuthCredentialExchanger class."""
@pytest.fixture
def base_exchanger(self):
return BaseAuthCredentialExchanger()
@pytest.fixture
def auth_scheme(self):
scheme = MagicMock(spec=AuthScheme)
scheme.type = "apiKey"
scheme.name = "x-api-key"
return scheme
def test_exchange_credential_not_implemented(
self, base_exchanger, auth_scheme
):
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.API_KEY, token="some-token"
)
with pytest.raises(NotImplementedError) as exc_info:
base_exchanger.exchange_credential(auth_scheme, auth_credential)
assert "Subclasses must implement exchange_credential." in str(
exc_info.value
)
def test_auth_credential_missing_error(self):
error_message = "Test missing credential"
error = AuthCredentialMissingError(error_message)
# assert error.message == error_message
assert str(error) == error_message

View File

@@ -0,0 +1,153 @@
# 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.
"""Tests for OAuth2CredentialExchanger."""
import copy
from unittest.mock import MagicMock
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import OAuth2Auth
from google.adk.auth.auth_schemes import AuthSchemeType
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
from google.adk.tools.openapi_tool.auth.credential_exchangers import OAuth2CredentialExchanger
from google.adk.tools.openapi_tool.auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError
import pytest
@pytest.fixture
def oauth2_exchanger():
return OAuth2CredentialExchanger()
@pytest.fixture
def auth_scheme():
openid_config = OpenIdConnectWithConfig(
type_=AuthSchemeType.openIdConnect,
authorization_endpoint="https://example.com/auth",
token_endpoint="https://example.com/token",
scopes=["openid", "profile"],
)
return openid_config
def test_check_scheme_credential_type_success(oauth2_exchanger, auth_scheme):
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="test_client",
client_secret="test_secret",
redirect_uri="http://localhost:8080",
),
)
# Check that the method does not raise an exception
oauth2_exchanger._check_scheme_credential_type(auth_scheme, auth_credential)
def test_check_scheme_credential_type_missing_credential(
oauth2_exchanger, auth_scheme
):
# Test case: auth_credential is None
with pytest.raises(ValueError) as exc_info:
oauth2_exchanger._check_scheme_credential_type(auth_scheme, None)
assert "auth_credential is empty" in str(exc_info.value)
def test_check_scheme_credential_type_invalid_scheme_type(
oauth2_exchanger, auth_scheme: OpenIdConnectWithConfig
):
"""Test case: Invalid AuthSchemeType."""
# Test case: Invalid AuthSchemeType
invalid_scheme = copy.deepcopy(auth_scheme)
invalid_scheme.type_ = AuthSchemeType.apiKey
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="test_client",
client_secret="test_secret",
redirect_uri="http://localhost:8080",
),
)
with pytest.raises(ValueError) as exc_info:
oauth2_exchanger._check_scheme_credential_type(
invalid_scheme, auth_credential
)
assert "Invalid security scheme" in str(exc_info.value)
def test_check_scheme_credential_type_missing_openid_connect(
oauth2_exchanger, auth_scheme
):
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
)
with pytest.raises(ValueError) as exc_info:
oauth2_exchanger._check_scheme_credential_type(auth_scheme, auth_credential)
assert "auth_credential is not configured with oauth2" in str(exc_info.value)
def test_generate_auth_token_success(
oauth2_exchanger, auth_scheme, monkeypatch
):
"""Test case: Successful generation of access token."""
# Test case: Successful generation of access token
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="test_client",
client_secret="test_secret",
redirect_uri="http://localhost:8080",
auth_response_uri="https://example.com/callback?code=test_code",
token={"access_token": "test_access_token"},
),
)
updated_credential = oauth2_exchanger.generate_auth_token(auth_credential)
assert updated_credential.auth_type == AuthCredentialTypes.HTTP
assert updated_credential.http.scheme == "bearer"
assert updated_credential.http.credentials.token == "test_access_token"
def test_exchange_credential_generate_auth_token(
oauth2_exchanger, auth_scheme, monkeypatch
):
"""Test exchange_credential when auth_response_uri is present."""
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.OAUTH2,
oauth2=OAuth2Auth(
client_id="test_client",
client_secret="test_secret",
redirect_uri="http://localhost:8080",
auth_response_uri="https://example.com/callback?code=test_code",
token={"access_token": "test_access_token"},
),
)
updated_credential = oauth2_exchanger.exchange_credential(
auth_scheme, auth_credential
)
assert updated_credential.auth_type == AuthCredentialTypes.HTTP
assert updated_credential.http.scheme == "bearer"
assert updated_credential.http.credentials.token == "test_access_token"
def test_exchange_credential_auth_missing(oauth2_exchanger, auth_scheme):
"""Test exchange_credential when auth_credential is missing."""
with pytest.raises(ValueError) as exc_info:
oauth2_exchanger.exchange_credential(auth_scheme, None)
assert "auth_credential is empty. Please create AuthCredential using" in str(
exc_info.value
)

View File

@@ -0,0 +1,196 @@
# 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.
"""Unit tests for the service account credential exchanger."""
from unittest.mock import MagicMock
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import ServiceAccount
from google.adk.auth.auth_credential import ServiceAccountCredential
from google.adk.auth.auth_schemes import AuthScheme
from google.adk.auth.auth_schemes import AuthSchemeType
from google.adk.tools.openapi_tool.auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError
from google.adk.tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger
import google.auth
import pytest
@pytest.fixture
def service_account_exchanger():
return ServiceAccountCredentialExchanger()
@pytest.fixture
def auth_scheme():
scheme = MagicMock(spec=AuthScheme)
scheme.type_ = AuthSchemeType.oauth2
scheme.description = "Google Service Account"
return scheme
def test_exchange_credential_success(
service_account_exchanger, auth_scheme, monkeypatch
):
"""Test successful exchange of service account credentials."""
mock_credentials = MagicMock()
mock_credentials.token = "mock_access_token"
# Mock the from_service_account_info method
mock_from_service_account_info = MagicMock(return_value=mock_credentials)
target_path = (
"google.adk.tools.openapi_tool.auth.credential_exchangers."
"service_account_exchanger.service_account.Credentials."
"from_service_account_info"
)
monkeypatch.setattr(
target_path,
mock_from_service_account_info,
)
# Mock the refresh method
mock_credentials.refresh = MagicMock()
# Create a valid AuthCredential with service account info
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
service_account=ServiceAccount(
service_account_credential=ServiceAccountCredential(
type_="service_account",
project_id="your_project_id",
private_key_id="your_private_key_id",
private_key="-----BEGIN PRIVATE KEY-----...",
client_email="...@....iam.gserviceaccount.com",
client_id="your_client_id",
auth_uri="https://accounts.google.com/o/oauth2/auth",
token_uri="https://oauth2.googleapis.com/token",
auth_provider_x509_cert_url=(
"https://www.googleapis.com/oauth2/v1/certs"
),
client_x509_cert_url=(
"https://www.googleapis.com/robot/v1/metadata/x509/..."
),
universe_domain="googleapis.com",
),
scopes=["https://www.googleapis.com/auth/cloud-platform"],
),
)
result = service_account_exchanger.exchange_credential(
auth_scheme, auth_credential
)
assert result.auth_type == AuthCredentialTypes.HTTP
assert result.http.scheme == "bearer"
assert result.http.credentials.token == "mock_access_token"
mock_from_service_account_info.assert_called_once()
mock_credentials.refresh.assert_called_once()
def test_exchange_credential_use_default_credential_success(
service_account_exchanger, auth_scheme, monkeypatch
):
"""Test successful exchange of service account credentials using default credential."""
mock_credentials = MagicMock()
mock_credentials.token = "mock_access_token"
mock_google_auth_default = MagicMock(
return_value=(mock_credentials, "test_project")
)
monkeypatch.setattr(google.auth, "default", mock_google_auth_default)
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
service_account=ServiceAccount(
use_default_credential=True,
scopes=["https://www.googleapis.com/auth/cloud-platform"],
),
)
result = service_account_exchanger.exchange_credential(
auth_scheme, auth_credential
)
assert result.auth_type == AuthCredentialTypes.HTTP
assert result.http.scheme == "bearer"
assert result.http.credentials.token == "mock_access_token"
mock_google_auth_default.assert_called_once()
mock_credentials.refresh.assert_called_once()
def test_exchange_credential_missing_auth_credential(
service_account_exchanger, auth_scheme
):
"""Test missing auth credential during exchange."""
with pytest.raises(AuthCredentialMissingError) as exc_info:
service_account_exchanger.exchange_credential(auth_scheme, None)
assert "Service account credentials are missing" in str(exc_info.value)
def test_exchange_credential_missing_service_account_info(
service_account_exchanger, auth_scheme
):
"""Test missing service account info during exchange."""
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
)
with pytest.raises(AuthCredentialMissingError) as exc_info:
service_account_exchanger.exchange_credential(auth_scheme, auth_credential)
assert "Service account credentials are missing" in str(exc_info.value)
def test_exchange_credential_exchange_failure(
service_account_exchanger, auth_scheme, monkeypatch
):
"""Test failure during service account token exchange."""
mock_from_service_account_info = MagicMock(
side_effect=Exception("Failed to load credentials")
)
target_path = (
"google.adk.tools.openapi_tool.auth.credential_exchangers."
"service_account_exchanger.service_account.Credentials."
"from_service_account_info"
)
monkeypatch.setattr(
target_path,
mock_from_service_account_info,
)
auth_credential = AuthCredential(
auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
service_account=ServiceAccount(
service_account_credential=ServiceAccountCredential(
type_="service_account",
project_id="your_project_id",
private_key_id="your_private_key_id",
private_key="-----BEGIN PRIVATE KEY-----...",
client_email="...@....iam.gserviceaccount.com",
client_id="your_client_id",
auth_uri="https://accounts.google.com/o/oauth2/auth",
token_uri="https://oauth2.googleapis.com/token",
auth_provider_x509_cert_url=(
"https://www.googleapis.com/oauth2/v1/certs"
),
client_x509_cert_url=(
"https://www.googleapis.com/robot/v1/metadata/x509/..."
),
universe_domain="googleapis.com",
),
scopes=["https://www.googleapis.com/auth/cloud-platform"],
),
)
with pytest.raises(AuthCredentialMissingError) as exc_info:
service_account_exchanger.exchange_credential(auth_scheme, auth_credential)
assert "Failed to exchange service account token" in str(exc_info.value)
mock_from_service_account_info.assert_called_once()