forked from databricks/databricks-sql-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_auth.py
More file actions
143 lines (113 loc) · 6.56 KB
/
test_auth.py
File metadata and controls
143 lines (113 loc) · 6.56 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import unittest
import pytest
from typing import Optional
from unittest.mock import patch
from databricks.sql.auth.auth import AccessTokenAuthProvider, BasicAuthProvider, AuthProvider, ExternalAuthProvider
from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
from databricks.sql.auth.oauth import OAuthManager
from databricks.sql.auth.authenticators import DatabricksOAuthProvider
from databricks.sql.auth.endpoint import CloudType, AwsOAuthEndpointCollection, AzureOAuthEndpointCollection
from databricks.sql.auth.authenticators import CredentialsProvider, HeaderFactory
from databricks.sql.experimental.oauth_persistence import OAuthPersistenceCache
class Auth(unittest.TestCase):
def test_access_token_provider(self):
access_token = "aBc2"
auth = AccessTokenAuthProvider(access_token=access_token)
http_request = {'myKey': 'myVal'}
auth.add_headers(http_request)
self.assertEqual(http_request['Authorization'], 'Bearer aBc2')
self.assertEqual(len(http_request.keys()), 2)
self.assertEqual(http_request['myKey'], 'myVal')
def test_basic_auth_provider(self):
username = "moderakh"
password = "Elevate Databricks 123!!!"
auth = BasicAuthProvider(username=username, password=password)
http_request = {'myKey': 'myVal'}
auth.add_headers(http_request)
self.assertEqual(http_request['Authorization'], 'Basic bW9kZXJha2g6RWxldmF0ZSBEYXRhYnJpY2tzIDEyMyEhIQ==')
self.assertEqual(len(http_request.keys()), 2)
self.assertEqual(http_request['myKey'], 'myVal')
def test_noop_auth_provider(self):
auth = AuthProvider()
http_request = {'myKey': 'myVal'}
auth.add_headers(http_request)
self.assertEqual(len(http_request.keys()), 1)
self.assertEqual(http_request['myKey'], 'myVal')
@patch.object(OAuthManager, "check_and_refresh_access_token")
@patch.object(OAuthManager, "get_tokens")
def test_oauth_auth_provider(self, mock_get_tokens, mock_check_and_refresh):
client_id = "mock-id"
scopes = ["offline_access", "sql"]
access_token = "mock_token"
refresh_token = "mock_refresh_token"
mock_get_tokens.return_value = (access_token, refresh_token)
mock_check_and_refresh.return_value = (access_token, refresh_token, False)
params = [(CloudType.AWS, "foo.cloud.databricks.com", AwsOAuthEndpointCollection, "offline_access sql"),
(CloudType.AZURE, "foo.1.azuredatabricks.net", AzureOAuthEndpointCollection,
f"{AzureOAuthEndpointCollection.DATATRICKS_AZURE_APP}/user_impersonation offline_access")]
for cloud_type, host, expected_endpoint_type, expected_scopes in params:
with self.subTest(cloud_type.value):
oauth_persistence = OAuthPersistenceCache()
auth_provider = DatabricksOAuthProvider(hostname=host,
oauth_persistence=oauth_persistence,
redirect_port_range=[8020],
client_id=client_id,
scopes=scopes)
self.assertIsInstance(auth_provider.oauth_manager.idp_endpoint, expected_endpoint_type)
self.assertEqual(auth_provider.oauth_manager.port_range, [8020])
self.assertEqual(auth_provider.oauth_manager.client_id, client_id)
self.assertEqual(oauth_persistence.read(host).refresh_token, refresh_token)
mock_get_tokens.assert_called_with(hostname=host, scope=expected_scopes)
headers = {}
auth_provider.add_headers(headers)
self.assertEqual(headers['Authorization'], f"Bearer {access_token}")
def test_external_provider(self):
class MyProvider(CredentialsProvider):
def auth_type(self) -> str:
return "mine"
def __call__(self, *args, **kwargs) -> HeaderFactory:
return lambda: {"foo": "bar"}
auth = ExternalAuthProvider(MyProvider())
http_request = {'myKey': 'myVal'}
auth.add_headers(http_request)
self.assertEqual(http_request['foo'], 'bar')
self.assertEqual(len(http_request.keys()), 2)
self.assertEqual(http_request['myKey'], 'myVal')
def test_get_python_sql_connector_auth_provider_access_token(self):
hostname = "moderakh-test.cloud.databricks.com"
kwargs = {'access_token': 'dpi123'}
auth_provider = get_python_sql_connector_auth_provider(hostname, **kwargs)
self.assertTrue(type(auth_provider).__name__, "AccessTokenAuthProvider")
headers = {}
auth_provider.add_headers(headers)
self.assertEqual(headers['Authorization'], 'Bearer dpi123')
def test_get_python_sql_connector_auth_provider_external(self):
class MyProvider(CredentialsProvider):
def auth_type(self) -> str:
return "mine"
def __call__(self, *args, **kwargs) -> HeaderFactory:
return lambda: {"foo": "bar"}
hostname = "moderakh-test.cloud.databricks.com"
kwargs = {'credentials_provider': MyProvider()}
auth_provider = get_python_sql_connector_auth_provider(hostname, **kwargs)
self.assertTrue(type(auth_provider).__name__, "ExternalAuthProvider")
headers = {}
auth_provider.add_headers(headers)
self.assertEqual(headers['foo'], 'bar')
def test_get_python_sql_connector_auth_provider_username_password(self):
username = "moderakh"
password = "Elevate Databricks 123!!!"
hostname = "moderakh-test.cloud.databricks.com"
kwargs = {'_username': username, '_password': password}
auth_provider = get_python_sql_connector_auth_provider(hostname, **kwargs)
self.assertTrue(type(auth_provider).__name__, "BasicAuthProvider")
headers = {}
auth_provider.add_headers(headers)
self.assertEqual(headers['Authorization'], 'Basic bW9kZXJha2g6RWxldmF0ZSBEYXRhYnJpY2tzIDEyMyEhIQ==')
def test_get_python_sql_connector_auth_provider_noop(self):
tls_client_cert_file = "fake.cert"
use_cert_as_auth = "abc"
hostname = "moderakh-test.cloud.databricks.com"
kwargs = {'_tls_client_cert_file': tls_client_cert_file, '_use_cert_as_auth': use_cert_as_auth}
auth_provider = get_python_sql_connector_auth_provider(hostname, **kwargs)
self.assertTrue(type(auth_provider).__name__, "CredentialProvider")