Skip to content

Commit ffce322

Browse files
ronpalronpal
committed
More suitable tests
1 parent 4c20cc2 commit ffce322

File tree

1 file changed

+257
-0
lines changed

1 file changed

+257
-0
lines changed

‎tests/tests_unit/test_credential_providers.py‎

Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from unittest.mock import Mock, patch
44

55
import pytest
6+
import requests
67
from oauthlib.oauth2 import InvalidClientIdError
78

89
from cognite.client.credentials import (
@@ -111,6 +112,111 @@ def test_access_token_generated(self, mock_public_client, expires_in):
111112
creds._refresh_access_token()
112113
assert "Authorization", "Bearer azure_token" == creds.authorization_header()
113114

115+
@patch("cognite.client.credentials.PublicClientApplication")
116+
@pytest.mark.parametrize(
117+
"authority_endpoint,authority_url,oauth_discovery_url,oidc_response,expected_endpoint,expected_error",
118+
[
119+
# MSAL authority has device_authorization_endpoint
120+
pytest.param(
121+
"https://login.microsoftonline.com/xyz/oauth2/v2.0/devicecode",
122+
"https://login.microsoftonline.com/xyz",
123+
None,
124+
None,
125+
"https://login.microsoftonline.com/xyz/oauth2/v2.0/devicecode",
126+
None,
127+
id="msal_authority_has_endpoint",
128+
),
129+
# Authority URL fallback (no MSAL endpoint)
130+
pytest.param(
131+
None,
132+
"https://login.microsoftonline.com/xyz",
133+
None,
134+
None,
135+
"https://login.microsoftonline.com/xyz/oauth2/v2.0/devicecode",
136+
None,
137+
id="authority_url_fallback",
138+
),
139+
# OIDC discovery with endpoint
140+
pytest.param(
141+
None,
142+
None,
143+
"https://auth0.example.com/oauth",
144+
{"device_authorization_endpoint": "https://auth0.example.com/oauth/device/code"},
145+
"https://auth0.example.com/oauth/device/code",
146+
None,
147+
id="oidc_discovery_with_endpoint",
148+
),
149+
# OIDC discovery missing endpoint
150+
pytest.param(
151+
None,
152+
None,
153+
"https://auth0.example.com/oauth",
154+
{"token_endpoint": "https://auth0.example.com/oauth/token"},
155+
None,
156+
"device_authorization_endpoint not found",
157+
id="oidc_discovery_missing_endpoint",
158+
),
159+
# OIDC discovery network error
160+
pytest.param(
161+
None,
162+
None,
163+
"https://auth0.example.com/oauth",
164+
requests.exceptions.RequestException("Network error"),
165+
None,
166+
"Error fetching device_authorization_endpoint from OIDC discovery",
167+
id="oidc_discovery_network_error",
168+
),
169+
],
170+
)
171+
def test_get_device_authorization_endpoint(
172+
self,
173+
mock_public_client,
174+
authority_endpoint,
175+
authority_url,
176+
oauth_discovery_url,
177+
oidc_response,
178+
expected_endpoint,
179+
expected_error,
180+
):
181+
"""Test _get_device_authorization_endpoint with various scenarios"""
182+
183+
# Setup authority mock
184+
mock_authority = Mock()
185+
mock_authority.device_authorization_endpoint = authority_endpoint
186+
mock_authority.instance = "instance" if authority_url else None
187+
mock_public_client().authority = mock_authority
188+
189+
# Setup OIDC discovery response
190+
if oidc_response is None:
191+
pass # No OIDC discovery
192+
elif isinstance(oidc_response, Exception):
193+
mock_public_client().http_client.get.side_effect = oidc_response
194+
else:
195+
mock_oidc_response = Mock()
196+
mock_oidc_response.json.return_value = oidc_response
197+
mock_public_client().http_client.get.return_value = mock_oidc_response
198+
199+
# Create OAuthDeviceCode instance (requires either authority_url or oauth_discovery_url)
200+
# For testing the method directly, we need at least one input
201+
if authority_url is None and oauth_discovery_url is None:
202+
# Skip this test case - can't create OAuthDeviceCode without inputs
203+
pytest.skip("Cannot test without authority_url or oauth_discovery_url")
204+
205+
creds = OAuthDeviceCode(
206+
authority_url=authority_url,
207+
oauth_discovery_url=oauth_discovery_url,
208+
client_id="test-client-id",
209+
scopes=["test-scope"],
210+
)
211+
212+
# Test the method
213+
if expected_error:
214+
with pytest.raises(CogniteAuthError, match=expected_error):
215+
creds._get_device_authorization_endpoint()
216+
else:
217+
result = creds._get_device_authorization_endpoint()
218+
assert result == expected_endpoint
219+
114220
@patch("cognite.client.credentials.PublicClientApplication")
115221
def test_entra_id_uses_authority_endpoint(self, mock_public_client):
116222
"""Test that Entra ID uses MSAL authority's device_authorization_endpoint when available"""
@@ -127,6 +233,7 @@ def test_entra_id_uses_authority_endpoint(self, mock_public_client):
127233
# Mock MSAL authority object with device_authorization_endpoint
128234
mock_authority = Mock()
129235
mock_authority.device_authorization_endpoint = "https://login.microsoftonline.com/xyz/oauth2/v2.0/devicecode"
236+
mock_authority.instance = "instance"
130237
mock_public_client().authority = mock_authority
131238
mock_public_client().http_client.post.return_value = mock_device_response
132239
mock_public_client().client.obtain_token_by_device_flow.return_value = {
@@ -157,6 +264,156 @@ def test_create_from_credential_provider(self, mock_public_client):
157264
assert isinstance(creds, OAuthDeviceCode)
158265
assert "Authorization", "Bearer azure_token" == creds.authorization_header()
159266

267+
@patch("cognite.client.credentials.PublicClientApplication")
268+
def test_oauth_discovery_url_device_flow(self, mock_public_client):
269+
"""Test that device code flow works with oauth_discovery_url (Auth0/Cognito)"""
270+
# Mock the OIDC discovery document response
271+
mock_oidc_response = Mock()
272+
mock_oidc_response.json.return_value = {
273+
"device_authorization_endpoint": "https://auth0.example.com/oauth/device/code",
274+
"token_endpoint": "https://auth0.example.com/oauth/token",
275+
}
276+
277+
# Mock the device code response
278+
mock_device_response = Mock()
279+
mock_device_response.json.return_value = {
280+
"user_code": "ABCD-EFGH",
281+
"device_code": "device123",
282+
"verification_uri": "https://auth0.example.com/activate",
283+
"expires_in": 900,
284+
"interval": 5,
285+
}
286+
287+
# Configure http_client to return different responses for GET (discovery) and POST (device code)
288+
def http_client_side_effect(url, **kwargs):
289+
if "openid-configuration" in url:
290+
return mock_oidc_response
291+
return mock_device_response
292+
293+
mock_public_client().http_client.get.side_effect = http_client_side_effect
294+
mock_public_client().http_client.post.side_effect = http_client_side_effect
295+
mock_public_client().client.obtain_token_by_device_flow.return_value = {
296+
"access_token": "auth0_token",
297+
"expires_in": 3600,
298+
}
299+
300+
# Mock authority to be None to simulate oauth_discovery_url behavior
301+
# But we need to mock instance for cache.add() call
302+
mock_authority = Mock()
303+
mock_authority.instance = None
304+
mock_authority.device_authorization_endpoint = None
305+
mock_public_client().authority = mock_authority
306+
307+
creds = OAuthDeviceCode(
308+
authority_url=None,
309+
oauth_discovery_url="https://auth0.example.com/oauth",
310+
client_id="auth0-client-id",
311+
scopes=["openid", "profile"],
312+
)
313+
creds._refresh_access_token()
314+
315+
# Verify OIDC discovery was fetched
316+
mock_public_client().http_client.get.assert_called_once_with(
317+
"https://auth0.example.com/oauth/.well-known/openid-configuration"
318+
)
319+
320+
# Verify the device authorization endpoint from discovery was called
321+
mock_public_client().http_client.post.assert_called_once()
322+
call_args = mock_public_client().http_client.post.call_args
323+
assert call_args[0][0] == "https://auth0.example.com/oauth/device/code"
324+
325+
assert "Authorization", "Bearer auth0_token" == creds.authorization_header()
326+
327+
@patch("cognite.client.credentials.PublicClientApplication")
328+
def test_device_code_response_normalized_response(self, mock_public_client):
329+
"""Test handling of MSAL NormalizedResponse (text fallback)"""
330+
mock_device_response = Mock()
331+
# Simulate NormalizedResponse that doesn't have .json() but has .text
332+
mock_device_response.json.side_effect = AttributeError("No json method")
333+
mock_device_response.text = '{"user_code": "ABCD", "device_code": "device123", "verification_uri": "https://example.com/activate", "expires_in": 900}'
334+
335+
mock_public_client().http_client.post.return_value = mock_device_response
336+
mock_public_client().client.obtain_token_by_device_flow.return_value = {
337+
"access_token": "token",
338+
"expires_in": 3600,
339+
}
340+
341+
# Mock authority to trigger fallback endpoint
342+
mock_authority = Mock()
343+
mock_authority.device_authorization_endpoint = None
344+
mock_public_client().authority = mock_authority
345+
346+
creds = OAuthDeviceCode(**self.DEFAULT_PROVIDER_ARGS)
347+
creds._refresh_access_token()
348+
349+
assert "Authorization", "Bearer token" == creds.authorization_header()
350+
351+
@patch("cognite.client.credentials.PublicClientApplication")
352+
def test_refresh_token_from_cache(self, mock_public_client):
353+
"""Test using refresh token from cache to get new access token"""
354+
import time
355+
356+
# Mock a valid refresh token in cache
357+
mock_refresh_token = {
358+
"secret": "refresh_token_secret",
359+
"expires_on": time.time() + 3600, # Valid for 1 hour
360+
}
361+
362+
mock_public_client().token_cache.search.return_value = [mock_refresh_token]
363+
mock_public_client().client.obtain_token_by_refresh_token.return_value = {
364+
"access_token": "new_access_token",
365+
"expires_in": 3600,
366+
}
367+
368+
creds = OAuthDeviceCode(**self.DEFAULT_PROVIDER_ARGS)
369+
creds._refresh_access_token()
370+
371+
# Verify refresh token was used
372+
mock_public_client().client.obtain_token_by_refresh_token.assert_called_once_with("refresh_token_secret")
373+
assert "Authorization", "Bearer new_access_token" == creds.authorization_header()
374+
375+
@patch("cognite.client.credentials.PublicClientApplication")
376+
def test_access_token_from_cache(self, mock_public_client):
377+
"""Test using valid access token from cache"""
378+
import time
379+
380+
# Mock no refresh token, but valid access token
381+
mock_public_client().token_cache.search.side_effect = [
382+
[], # No refresh tokens
383+
[ # Access tokens
384+
{
385+
"secret": "cached_access_token",
386+
"expires_on": str(int(time.time() + 3600)), # Valid for 1 hour
387+
}
388+
],
389+
]
390+
391+
creds = OAuthDeviceCode(**self.DEFAULT_PROVIDER_ARGS)
392+
creds._refresh_access_token()
393+
394+
# Verify device flow was NOT triggered
395+
mock_public_client().http_client.post.assert_not_called()
396+
assert "Authorization", "Bearer cached_access_token" == creds.authorization_header()
397+
398+
@patch("cognite.client.credentials.PublicClientApplication")
399+
def test_device_flow_error_response(self, mock_public_client):
400+
"""Test handling of error response from device authorization endpoint"""
401+
mock_device_response = Mock()
402+
mock_device_response.json.return_value = {
403+
"error": "invalid_client",
404+
"error_description": "Invalid client credentials",
405+
}
406+
407+
mock_public_client().http_client.post.return_value = mock_device_response
408+
mock_authority = Mock()
409+
mock_authority.device_authorization_endpoint = None
410+
mock_public_client().authority = mock_authority
411+
412+
creds = OAuthDeviceCode(**self.DEFAULT_PROVIDER_ARGS)
413+
414+
with pytest.raises(CogniteAuthError, match=r"Error initiating device flow.*invalid_client"):
415+
creds._refresh_access_token()
416+
160417

161418
class TestOAuthInteractive:
162419
DEFAULT_PROVIDER_ARGS: ClassVar = MappingProxyType(

0 commit comments

Comments
 (0)