Skip to content

Commit 49b65d3

Browse files
parthban-dbparthban-db
committed
fix tests
1 parent d4b95e3 commit 49b65d3

File tree

6 files changed

+23
-23
lines changed

6 files changed

+23
-23
lines changed

‎tests/test_config.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ def test_config_supports_legacy_credentials_provider():
3939
],
4040
)
4141
def test_config_host_url_format_check(mocker, host, expected):
42-
mocker.patch("databricks.sdk.config.Config.init_auth")
42+
mocker.patch("databricks.sdk.databricks.config.Config.init_auth")
4343
assert Config(host=host).host == expected
4444

4545

@@ -113,7 +113,7 @@ def test_config_copy_deep_copies_user_agent_other_info(config):
113113

114114
def test_config_deep_copy(monkeypatch, mocker, tmp_path):
115115
mocker.patch(
116-
"databricks.sdk.credentials_provider.CliTokenSource.refresh",
116+
"databricks.sdk.databricks.credentials_provider.CliTokenSource.refresh",
117117
return_value=Token(
118118
access_token="token",
119119
token_type="Bearer",

‎tests/test_core.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ def test_databricks_cli_credential_provider_installed_legacy(config, monkeypatch
166166

167167
def test_databricks_cli_credential_provider_installed_new(config, monkeypatch, tmp_path, mocker):
168168
get_mock = mocker.patch(
169-
"databricks.sdk.credentials_provider.CliTokenSource.refresh",
169+
"databricks.sdk.databricks.credentials_provider.CliTokenSource.refresh",
170170
return_value=Token(
171171
access_token="token",
172172
token_type="Bearer",

‎tests/test_credentials_provider.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ def test_external_browser_refresh_success(mocker):
2626

2727
# Inject the mock implementations.
2828
mocker.patch(
29-
"databricks.sdk.credentials_provider.TokenCache",
29+
"databricks.sdk.databricks.credentials_provider.TokenCache",
3030
return_value=mock_token_cache,
3131
)
3232

@@ -66,11 +66,11 @@ def test_external_browser_refresh_failure_new_oauth_flow(mocker):
6666

6767
# Inject the mock implementations.
6868
mocker.patch(
69-
"databricks.sdk.credentials_provider.TokenCache",
69+
"databricks.sdk.databricks.credentials_provider.TokenCache",
7070
return_value=mock_token_cache,
7171
)
7272
mocker.patch(
73-
"databricks.sdk.credentials_provider.OAuthClient",
73+
"databricks.sdk.databricks.credentials_provider.OAuthClient",
7474
return_value=mock_oauth_client,
7575
)
7676

@@ -112,11 +112,11 @@ def test_external_browser_no_cached_credentials(mocker):
112112

113113
# Inject the mock implementations.
114114
mocker.patch(
115-
"databricks.sdk.credentials_provider.TokenCache",
115+
"databricks.sdk.databricks.credentials_provider.TokenCache",
116116
return_value=mock_token_cache,
117117
)
118118
mocker.patch(
119-
"databricks.sdk.credentials_provider.OAuthClient",
119+
"databricks.sdk.databricks.credentials_provider.OAuthClient",
120120
return_value=mock_oauth_client,
121121
)
122122

@@ -150,11 +150,11 @@ def test_external_browser_consent_fails(mocker):
150150

151151
# Inject the mock implementations.
152152
mocker.patch(
153-
"databricks.sdk.credentials_provider.TokenCache",
153+
"databricks.sdk.databricks.credentials_provider.TokenCache",
154154
return_value=mock_token_cache,
155155
)
156156
mocker.patch(
157-
"databricks.sdk.credentials_provider.OAuthClient",
157+
"databricks.sdk.databricks.credentials_provider.OAuthClient",
158158
return_value=mock_oauth_client,
159159
)
160160

‎tests/test_data_plane.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def test_endpoint_token_source_get_token(config):
2222
config.host, success_callable(cp_token), "authDetails", disable_async=True
2323
)
2424

25-
with patch("databricks.sdk.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
25+
with patch("databricks.sdk.databricks.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
2626
token_source.token()
2727

2828
retrieve_token.assert_called_once()
@@ -42,7 +42,7 @@ def test_endpoint_token_source_get_token(config):
4242
def test_token_source_get_token_not_existing(config):
4343
token_source = data_plane.DataPlaneTokenSource(config.host, success_callable(cp_token), disable_async=True)
4444

45-
with patch("databricks.sdk.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
45+
with patch("databricks.sdk.databricks.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
4646
result_token = token_source.token(endpoint="endpoint", auth_details="authDetails")
4747

4848
retrieve_token.assert_called_once()
@@ -64,7 +64,7 @@ def test_token_source_get_token_existing(config):
6464
token_source = data_plane.DataPlaneTokenSource(config.host, success_callable(cp_token), disable_async=True)
6565
token_source._token_sources["endpoint:authDetails"] = MockEndpointTokenSource(another_token)
6666

67-
with patch("databricks.sdk.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
67+
with patch("databricks.sdk.databricks.oauth.retrieve_token", return_value=dp_token) as retrieve_token:
6868
result_token = token_source.token(endpoint="endpoint", auth_details="authDetails")
6969

7070
retrieve_token.assert_not_called()

‎tests/test_dbutils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,7 @@ def test_any_proxy(dbutils_proxy):
251251

252252

253253
def test_secrets_get_and_redacting_logs(dbutils, mocker):
254-
inner = mocker.patch("databricks.sdk.core.ApiClient.do", return_value={"value": "aGVsbG8="})
254+
inner = mocker.patch("databricks.sdk.databricks.core.ApiClient.do", return_value={"value": "aGVsbG8="})
255255

256256
value = dbutils.secrets.get("foo", "bar")
257257

‎tests/test_model_serving_auth.py‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -62,10 +62,10 @@ def test_model_serving_auth(env_values, del_env_values, oauth_file_name, monkeyp
6262

6363
# patch mlflow to read the file from the test directory
6464
monkeypatch.setattr(
65-
"databricks.sdk.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
65+
"databricks.sdk.databricks.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
6666
oauth_file_name,
6767
)
68-
mocker.patch("databricks.sdk.config.Config._known_file_config_loader")
68+
mocker.patch("databricks.sdk.databricks.config.Config._known_file_config_loader")
6969

7070
cfg = Config()
7171
assert cfg.auth_type == "model-serving"
@@ -107,7 +107,7 @@ def test_model_serving_auth_errors(env_values, oauth_file_name, monkeypatch):
107107
for env_name, env_value in env_values:
108108
monkeypatch.setenv(env_name, env_value)
109109
monkeypatch.setattr(
110-
"databricks.sdk.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
110+
"databricks.sdk.databricks.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
111111
oauth_file_name,
112112
)
113113

@@ -121,10 +121,10 @@ def test_model_serving_auth_refresh(monkeypatch, mocker):
121121

122122
# patch mlflow to read the file from the test directory
123123
monkeypatch.setattr(
124-
"databricks.sdk.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
124+
"databricks.sdk.databricks.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
125125
"tests/testdata/model-serving-test-token",
126126
)
127-
mocker.patch("databricks.sdk.config.Config._known_file_config_loader")
127+
mocker.patch("databricks.databricks.sdk.config.Config._known_file_config_loader")
128128

129129
cfg = Config()
130130
assert cfg.auth_type == "model-serving"
@@ -135,12 +135,12 @@ def test_model_serving_auth_refresh(monkeypatch, mocker):
135135
assert headers.get("Authorization") == "Bearer databricks_sdk_unit_test_token" # Token defined in the test file
136136
# Simulate refreshing the token by patching to to a new file
137137
monkeypatch.setattr(
138-
"databricks.sdk.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
138+
"databricks.sdk.databricks.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
139139
"tests/testdata/model-serving-test-token-v2",
140140
)
141141

142142
monkeypatch.setattr(
143-
"databricks.sdk.credentials_provider.time.time",
143+
"databricks.sdk.databricks.credentials_provider.time.time",
144144
lambda: current_time + 10,
145145
)
146146

@@ -151,7 +151,7 @@ def test_model_serving_auth_refresh(monkeypatch, mocker):
151151

152152
# Expiry is 300 seconds so this should force an expiry and re read from the new file path
153153
monkeypatch.setattr(
154-
"databricks.sdk.credentials_provider.time.time",
154+
"databricks.sdk.databricks.credentials_provider.time.time",
155155
lambda: current_time + 600,
156156
)
157157

@@ -171,7 +171,7 @@ def test_agent_user_credentials(monkeypatch, mocker):
171171
monkeypatch.setenv("IS_IN_DB_MODEL_SERVING_ENV", "true")
172172
monkeypatch.setenv("DB_MODEL_SERVING_HOST_URL", "x")
173173
monkeypatch.setattr(
174-
"databricks.sdk.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
174+
"databricks.sdk.databricks.credentials_provider.ModelServingAuthProvider._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH",
175175
"tests/testdata/model-serving-test-token",
176176
)
177177

0 commit comments

Comments
 (0)