8000 chore: Add credential exchanger registry (Experimentals) · DavidSchmidt00/adk-python@55201cb · GitHub
[go: up one dir, main page]
Skip to content

Commit 55201cb

Browse files
seanzhougooglecopybara-github
seanzhougoogle
authored and
copybara-github
committed
chore: Add credential exchanger registry (Experimentals)
PiperOrigin-RevId: 772713412
1 parent 0a96253 commit 55201cb

File tree

2 files changed

+300
-0
lines changed

2 files changed

+300
-0
lines changed

‎src/google/adk/auth/exchanger/credential_exchanger_registry.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Credential exchanger registry."""
16+
17+
from __future__ import annotations
18+
19+
from typing import Dict
20+
from typing import Optional
21+
22+
from ...utils.feature_decorator import experimental
23+
from ..auth_credential import AuthCredentialTypes
24+
from .base_credential_exchanger import BaseCredentialExchanger
25+
26+
27+
@experimental
28+
class CredentialExchangerRegistry:
29+
"""Registry for credential exchanger instances."""
30+
31+
def __init__(self):
32+
self._exchangers: Dict[AuthCredentialTypes, BaseCredentialExchanger] = {}
33+
34+
def register(
35+
self,
36+
credential_type: AuthCredentialTypes,
37+
exchanger_instance: BaseCredentialExchanger,
38+
) -> None:
39+
"""Register an exchanger instance for a credential type.
40+
41+
Args:
42+
credential_type: The credential type to register for.
43+
exchanger_instance: The exchanger instance to register.
44+
"""
45+
self._exchangers[credential_type] = exchanger_instance
46+
47+
def get_exchanger(
48+
self, credential_type: AuthCredentialTypes
49+
) -> Optional[BaseCredentialExchanger]:
50+
"""Get the exchanger instance for a credential type.
51+
52+
Args:
53+
credential_type: The credential type to get exchanger for.
54+
55+
Returns:
56+
The exchanger instance if registered, None otherwise.
57+
"""
58+
return self._exchangers.get(credential_type)

‎tests/unittests/auth/exchanger/test_credential_exchanger_registry.py

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Unit tests for the CredentialExchangerRegistry."""
16+
17+
from typing import Optional
18+
from unittest.mock import MagicMock
19+
20+
from google.adk.auth.auth_credential import AuthCredential
21+
from google.adk.auth.auth_credential import AuthCredentialTypes
22+
from google.adk.auth.auth_schemes import AuthScheme
23+
from google.adk.auth.exchanger.base_credential_exchanger import BaseCredentialExchanger
24+
from google.adk.auth.exchanger.credential_exchanger_registry import CredentialExchangerRegistry
25+
import pytest
26+
27+
28+
class MockCredentialExchanger(BaseCredentialExchanger):
29+
"""Mock credential exchanger for testing."""
30+
31+
def __init__(self, exchange_result: Optional[AuthCredential] = None):
32+
self.exchange_result = exchange_result or AuthCredential(
33+
auth_type=AuthCredentialTypes.HTTP
34+
)
35+
36+
def exchange(
37+
self,
38+
auth_credential: AuthCredential,
39+
auth_scheme: Optional[AuthScheme] = None,
40+
) -> AuthCredential:
41+
"""Mock exchange method."""
42+
return self.exchange_result
43+
44+
45+
class TestCredentialExchangerRegistry:
46+
"""Test cases for CredentialExchangerRegistry."""
47+
48+
def test_initialization(self):
49+
"""Test that the registry initializes with an empty exchangers dictionary."""
50+
registry = CredentialExchangerRegistry()
51+
52+
# Access the private attribute for testing
53+
assert hasattr(registry, '_exchangers')
54+
assert isinstance(registry._exchangers, dict)
55+
assert len(registry._exchangers) == 0
56+
57+
def test_register_single_exchanger(self):
58+
"""Test registering a single exchanger."""
59+
registry = CredentialExchangerRegistry()
60+
mock_exchanger = MockCredentialExchanger()
61+
62+
registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
63+
64+
# Verify the exchanger was registered
65+
retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.API_KEY)
66+
assert retrieved_exchanger is mock_exchanger
67+
68+
def test_register_multiple_exchangers(self):
69+
"""Test registering multiple exchangers for different credential types."""
70+
registry = CredentialExchangerRegistry()
71+
72+
api_key_exchanger = MockCredentialExchanger()
73+
oauth2_exchanger = MockCredentialExchanger()
74+
service_account_exchanger = MockCredentialExchanger()
75+
76+
registry.register(AuthCredentialTypes.API_KEY, api_key_exchanger)
77+
registry.register(AuthCredentialTypes.OAUTH2, oauth2_exchanger)
78+
registry.register(
79+
AuthCredentialTypes.SERVICE_ACCOUNT, service_account_exchanger
80+
)
81+
82+
# Verify all exchangers were registered correctly
83+
assert (
84+
registry.get_exchanger(AuthCredentialTypes.API_KEY) is api_key_exchanger
85+
)
86+
assert (
87+
registry.get_exchanger(AuthCredentialTypes.OAUTH2) is oauth2_exchanger
88+
)
89+
assert (
90+
registry.get_exchanger(AuthCredentialTypes.SERVICE_ACCOUNT)
91+
is service_account_exchanger
92+
)
93+
94+
def test_register_overwrites_existing_exchanger(self):
95+
"""Test that registering an exchanger for an existing type overwrites the previous one."""
96+
registry = CredentialExchangerRegistry()
97+
98+
first_exchanger = MockCredentialExchanger()
99+
second_exchanger = MockCredentialExchanger()
100+
101+
# Register first exchanger
102+
registry.register(AuthCredentialTypes.API_KEY, first_exchanger)
103+
assert (
104+
registry.get_exchanger(AuthCredentialTypes.API_KEY) is first_exchanger
105+
)
106+
107+
# Register second exchanger for the same type
108+
registry.register(AuthCredentialTypes.API_KEY, second_exchanger)
109+
assert (
110+
registry.get_exchanger(AuthCredentialTypes.API_KEY) is second_exchanger
111+
)
112+
assert (
113+
registry.get_exchanger(AuthCredentialTypes.API_KEY)
114+
is not first_exchanger
115+
)
116+
117+
def test_get_exchanger_returns_correct_instance(self):
118+
"""Test that get_exchanger returns the correct exchanger instance."""
119+
registry = CredentialExchangerRegistry()
120+
mock_exchanger = MockCredentialExchanger()
121+
122+
registry.register(AuthCredentialTypes.HTTP, mock_exchanger)
123+
124+
retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.HTTP)
125+
assert retrieved_exchanger is mock_exchanger
126+
assert isinstance(retrieved_exchanger, BaseCredentialExchanger)
127+
128+
def test_get_exchanger_nonexistent_type_returns_none(self):
129+
"""Test that get_exchanger returns None for non-existent credential types."""
130+
registry = CredentialExchangerRegistry()
131+
132+
# Try to get an exchanger that was never registered
133+
result = registry.get_exchanger(AuthCredentialTypes.OAUTH2)
134+
assert result is None
135+
136+
def test_get_exchanger_after_registration_and_removal(self):
137+
"""Test behavior when an exchanger is registered and then the registry is cleared indirectly."""
138+
registry = CredentialExchangerRegistry()
139+
mock_exchanger = MockCredentialExchanger()
140+
141+
# Register exchanger
142+
registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
143< 10000 /td>+
assert registry.get_exchanger(AuthCredentialTypes.API_KEY) is mock_exchanger
144+
145+
# Clear the internal dictionary (simulating some edge case)
146+
registry._exchangers.clear()
147+
assert registry.get_exchanger(AuthCredentialTypes.API_KEY) is None
148+
149+
def test_register_with_all_credential_types(self):
150+
"""Test registering exchangers for all available credential types."""
151+
registry = CredentialExchangerRegistry()
152+
153+
exchangers = {}
154+
credential_types = [
155+
AuthCredentialTypes.API_KEY,
156+
AuthCredentialTypes.HTTP,
157+
AuthCredentialTypes.OAUTH2,
158+
AuthCredentialTypes.OPEN_ID_CONNECT,
159+
AuthCredentialTypes.SERVICE_ACCOUNT,
160+
]
161+
162+
# Register an exchanger for each credential type
163+
for cred_type in credential_types:
164+
exchanger = MockCredentialExchanger()
165+
exchangers[cred_type] = exchanger
166+
registry.register(cred_type, exchanger)
167+
168+
# Verify all exchangers can be retrieved
169+
for cred_type in credential_types:
170+
retrieved_exchanger = registry.get_exchanger(cred_type)
171+
assert retrieved_exchanger is exchangers[cred_type]
172+
173+
def test_register_with_mock_exchanger_using_magicmock(self):
174+
"""Test registering with a MagicMock exchanger."""
175+
registry = CredentialExchangerRegistry()
176+
mock_exchanger = MagicMock(spec=BaseCredentialExchanger)
177+
178+
registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
179+
180+
retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.API_KEY)
181+
assert retrieved_exchanger is mock_exchanger
182+
183+
def test_registry_isolation(self):
184+
"""Test that different registry instances are isolated from each other."""
185+
registry1 = CredentialExchangerRegistry()
186+
registry2 = CredentialExchangerRegistry()
187+
188+
exchanger1 = MockCredentialExchanger()
189+
exchanger2 = MockCredentialExchanger()
190+
191+
# Register different exchangers in different registry instances
192+
registry1.register(AuthCredentialTypes.API_KEY, exchanger1)
193+
registry2.register(AuthCredentialTypes.API_KEY, exchanger2)
194+
195+
# Verify isolation
196+
assert registry1.get_exchanger(AuthCredentialTypes.API_KEY) is exchanger1
197+
assert registry2.get_exchanger(AuthCredentialTypes.API_KEY) is exchanger2
198+
assert (
199+
registry1.get_exchanger(AuthCredentialTypes.API_KEY) is not exchanger2
200+
)
201+
assert (
202+
registry2.get_exchanger(AuthCredentialTypes.API_KEY) is not exchanger1
203+
)
204+
205+
def test_exchanger_functionality_through_registry(self):
206+
"""Test that exchangers registered in the registry function correctly."""
207+
registry = CredentialExchangerRegistry()
208+
209+
# Create a mock exchanger with specific return value
210+
expected_result = AuthCredential(auth_type=AuthCredentialTypes.HTTP)
211+
mock_exchanger = MockCredentialExchanger(exchange_result=expected_result)
212+
213+
registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
214+
215+
# Get the exchanger and test its functionality
216+
retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.API_KEY)
217+
input_credential = AuthCredential(auth_type=AuthCredentialTypes.API_KEY)
218+
219+
result = retrieved_exchanger.exchange(input_credential)
220+
assert result is expected_result
221+
222+
def test_register_none_exchanger(self):
223+
"""Test that registering None as an exchanger works (edge case)."""
224+
registry = CredentialExchangerRegistry()
225+
226+
# This should work but return None when retrieved
227+
registry.register(AuthCredentialTypes.API_KEY, None)
228+
229+
result = registry.get_exchanger(AuthCredentialTypes.API_KEY)
230+
assert result is None
231+
232+
def test_internal_dictionar 6D57 y_structure(self):
233+
"""Test the internal structure of the registry."""
234+
registry = CredentialExchangerRegistry()
235+
mock_exchanger = MockCredentialExchanger()
236+
237+
registry.register(AuthCredentialTypes.OAUTH2, mock_exchanger)
238+
239+
# Verify internal dictionary structure
240+
assert AuthCredentialTypes.OAUTH2 in registry._exchangers
241+
assert registry._exchangers[AuthCredentialTypes.OAUTH2] is mock_exchanger
242+
assert len(registry._exchangers) == 1

0 commit comments

Comments
 (0)
0