mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: add audience and prompt as configurable for OAuth flows
Merge https://github.com/google/adk-python/pull/2738 Some OAuth servers require audience such as [Jira](https://developer.atlassian.com/cloud/confluence/oauth-2-3lo-apps/). This change allows the configuration of audience and prompt (if it needs to be changed) and adds some tests. This is for issue [2755](https://github.com/google/adk-python/issues/2755). Resolves #2755 Happy to provide changes/updates if needed. COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2738 from mescanne:oauth-audience-prompt 87ce1100792d9156ada2a004bcfaf2fe5fc69602 PiperOrigin-RevId: 802850034
This commit is contained in:
committed by
Copybara-Service
parent
a1679dae3f
commit
edda922791
@@ -79,6 +79,7 @@ class OAuth2Auth(BaseModelWithConfig):
|
||||
refresh_token: Optional[str] = None
|
||||
expires_at: Optional[int] = None
|
||||
expires_in: Optional[int] = None
|
||||
audience: Optional[str] = None
|
||||
|
||||
|
||||
class ServiceAccountCredential(BaseModelWithConfig):
|
||||
|
||||
@@ -188,9 +188,16 @@ class AuthHandler:
|
||||
scope=" ".join(scopes),
|
||||
redirect_uri=auth_credential.oauth2.redirect_uri,
|
||||
)
|
||||
params = {
|
||||
"access_type": "offline",
|
||||
"prompt": "consent",
|
||||
}
|
||||
if auth_credential.oauth2.audience:
|
||||
params["audience"] = auth_credential.oauth2.audience
|
||||
uri, state = client.create_authorization_url(
|
||||
url=authorization_endpoint, access_type="offline", prompt="consent"
|
||||
url=authorization_endpoint, **params
|
||||
)
|
||||
|
||||
exchanged_auth_credential = auth_credential.model_copy(deep=True)
|
||||
exchanged_auth_credential.oauth2.auth_uri = uri
|
||||
exchanged_auth_credential.oauth2.state = state
|
||||
|
||||
@@ -61,7 +61,10 @@ class MockOAuth2Session:
|
||||
self.state = state
|
||||
|
||||
def create_authorization_url(self, url, **kwargs):
|
||||
return f"{url}?client_id={self.client_id}&scope={self.scope}", "mock_state"
|
||||
params = f"client_id={self.client_id}&scope={self.scope}"
|
||||
if kwargs.get("audience"):
|
||||
params += f"&audience={kwargs.get('audience')}"
|
||||
return f"{url}?{params}", "mock_state"
|
||||
|
||||
def fetch_token(
|
||||
self,
|
||||
@@ -225,8 +228,27 @@ class TestGenerateAuthUri:
|
||||
"https://example.com/oauth2/authorize"
|
||||
)
|
||||
assert "client_id=mock_client_id" in result.oauth2.auth_uri
|
||||
assert "audience" not in result.oauth2.auth_uri
|
||||
assert result.oauth2.state == "mock_state"
|
||||
|
||||
@patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
|
||||
def test_generate_auth_uri_with_audience_and_prompt(
|
||||
self, openid_auth_scheme, oauth2_credentials
|
||||
):
|
||||
"""Test generating an auth URI with audience and prompt."""
|
||||
oauth2_credentials.oauth2.audience = "test_audience"
|
||||
exchanged = oauth2_credentials.model_copy(deep=True)
|
||||
|
||||
config = AuthConfig(
|
||||
auth_scheme=openid_auth_scheme,
|
||||
raw_auth_credential=oauth2_credentials,
|
||||
exchanged_auth_credential=exchanged,
|
||||
)
|
||||
handler = AuthHandler(config)
|
||||
result = handler.generate_auth_uri()
|
||||
|
||||
assert "audience=test_audience" in result.oauth2.auth_uri
|
||||
|
||||
@patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
|
||||
def test_generate_auth_uri_openid(
|
||||
self, openid_auth_scheme, oauth2_credentials
|
||||
|
||||
Reference in New Issue
Block a user