def test_databricks_params_no_verify(get_config):
get_config.return_value = DatabricksConfig.from_password("host",
"user",
"pass",
insecure=True)
params = databricks_utils.get_databricks_host_creds()
assert params.ignore_tls_verification
def test_databricks_http_request_integration(get_config, request):
"""Confirms that the databricks http request params can in fact be used as an HTTP request"""
def confirm_request_params(*args, **kwargs):
headers = dict(_DEFAULT_HEADERS)
headers["Authorization"] = "Basic dXNlcjpwYXNz"
assert args == ("PUT", "host/clusters/list")
assert kwargs == {
"headers": headers,
"verify": True,
"json": {"a": "b"},
"timeout": 120,
}
http_response = mock.MagicMock()
http_response.status_code = 200
http_response.text = '{"OK": "woo"}'
return http_response
request.side_effect = confirm_request_params
get_config.return_value = DatabricksConfig.from_password("host", "user", "pass", insecure=False)
response = DatabricksJobRunner(databricks_profile_uri=None)._databricks_api_request(
"/clusters/list", "PUT", json={"a": "b"}
)
assert json.loads(response.text) == {"OK": "woo"}
get_config.reset_mock()
response = DatabricksJobRunner(
databricks_profile_uri=construct_db_uri_from_profile("my-profile")
)._databricks_api_request("/clusters/list", "PUT", json={"a": "b"})
assert json.loads(response.text) == {"OK": "woo"}
assert get_config.call_count == 0
def test_databricks_params_throws_errors(ProfileConfigProvider):
# No hostname
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = DatabricksConfig.from_password(
None, "user", "pass", insecure=True)
ProfileConfigProvider.return_value = mock_provider
with pytest.raises(Exception, match="You haven't configured the CLI yet"):
databricks_utils.get_databricks_host_creds()
# No authentication
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = DatabricksConfig.from_password(
"host", None, None, insecure=True)
ProfileConfigProvider.return_value = mock_provider
with pytest.raises(Exception, match="You haven't configured the CLI yet"):
databricks_utils.get_databricks_host_creds()
def test_databricks_params_custom_profile(ProfileConfigProvider):
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = DatabricksConfig.from_password(
"host", "user", "pass", insecure=True
)
ProfileConfigProvider.return_value = mock_provider
params = databricks_utils.get_databricks_host_creds(construct_db_uri_from_profile("profile"))
assert params.ignore_tls_verification
ProfileConfigProvider.assert_called_with("profile")
def test_databricks_params_user_password(get_config):
get_config.return_value = DatabricksConfig.from_password("host",
"user",
"pass",
insecure=False)
params = databricks_utils.get_databricks_host_creds()
assert params.host == "host"
assert params.username == "user"
assert params.password == "pass"
def test_update_and_persist_config_case_sensitive():
config = DatabricksConfig.from_token(TEST_HOST, TEST_TOKEN)
update_and_persist_config(TEST_PROFILE, config)
config_2 = DatabricksConfig.from_password(TEST_HOST, TEST_USER, TEST_PASSWORD)
update_and_persist_config(TEST_PROFILE.upper(), config_2)
config = get_config_for_profile(TEST_PROFILE)
assert config.is_valid_with_token
assert not config.is_valid_with_password
config_2 = get_config_for_profile(TEST_PROFILE.upper())
assert config_2.is_valid_with_password
assert not config_2.is_valid_with_token
def _configure_cli_password(profile, insecure):
config = get_config_for_profile(profile)
if config.password:
default_password = '******' * len(config.password)
else:
default_password = None
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
username = click.prompt(PROMPT_USERNAME, default=config.username)
password = click.prompt(PROMPT_PASSWORD, default=default_password, hide_input=True,
confirmation_prompt=True)
if password == default_password:
password = config.password
new_config = DatabricksConfig.from_password(host, username, password, insecure)
update_and_persist_config(profile, new_config)
def test_docker_databricks_tracking_cmd_and_envs(ProfileConfigProvider):
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = DatabricksConfig.from_password(
"host", "user", "pass", insecure=True)
ProfileConfigProvider.return_value = mock_provider
cmds, envs = mlflow.projects.docker.get_docker_tracking_cmd_and_envs(
"databricks://some-profile")
assert envs == {
"DATABRICKS_HOST": "host",
"DATABRICKS_USERNAME": "******",
"DATABRICKS_PASSWORD": "******",
"DATABRICKS_INSECURE": "True",
mlflow.tracking._TRACKING_URI_ENV_VAR: "databricks",
}
assert cmds == []
def test_update_and_persist_config():
config = DatabricksConfig.from_token(TEST_HOST, TEST_TOKEN)
update_and_persist_config(DEFAULT_SECTION, config)
config = get_config_for_profile(DEFAULT_SECTION)
assert config.is_valid_with_token
assert config.host == TEST_HOST
assert config.token == TEST_TOKEN
# Overwrite conf for same section.
config = DatabricksConfig.from_password(TEST_HOST, TEST_USER, TEST_PASSWORD)
update_and_persist_config(DEFAULT_SECTION, config)
config = get_config_for_profile(DEFAULT_SECTION)
assert config.is_valid_with_password
assert not config.is_valid_with_token
assert config.host == TEST_HOST
assert config.username == TEST_USER
assert config.password == TEST_PASSWORD
def _configure_cli_password(profile, insecure, host, jobs_api_version):
config = ProfileConfigProvider(
profile).get_config() or DatabricksConfig.empty()
if config.password:
default_password = '******' * len(config.password)
else:
default_password = None
if not host:
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
username = click.prompt(PROMPT_USERNAME, default=config.username)
password = click.prompt(PROMPT_PASSWORD,
default=default_password,
hide_input=True,
confirmation_prompt=True)
if password == default_password:
password = config.password
new_config = DatabricksConfig.from_password(host, username, password,
insecure, jobs_api_version)
update_and_persist_config(profile, new_config)
def test_docker_project_tracking_uri_propagation(
ProfileConfigProvider, tmpdir, tracking_uri, expected_command_segment, docker_example_base_image
): # pylint: disable=unused-argument
mock_provider = mock.MagicMock()
mock_provider.get_config.return_value = DatabricksConfig.from_password(
"host", "user", "pass", insecure=True
)
ProfileConfigProvider.return_value = mock_provider
# Create and mock local tracking directory
local_tracking_dir = os.path.join(tmpdir.strpath, "mlruns")
if tracking_uri is None:
tracking_uri = local_tracking_dir
old_uri = mlflow.get_tracking_uri()
try:
mlflow.set_tracking_uri(tracking_uri)
with mock.patch("mlflow.tracking._tracking_service.utils._get_store") as _get_store_mock:
_get_store_mock.return_value = file_store.FileStore(local_tracking_dir)
mlflow.projects.run(
TEST_DOCKER_PROJECT_DIR, experiment_id=file_store.FileStore.DEFAULT_EXPERIMENT_ID
)
finally:
mlflow.set_tracking_uri(old_uri)
def test_is_valid(self):
config = DatabricksConfig.from_password(TEST_HOST, TEST_USER,
TEST_PASSWORD)
assert config.is_valid
config = DatabricksConfig.from_token(TEST_HOST, TEST_TOKEN)
assert config.is_valid
def test_is_valid_with_password(self):
config = DatabricksConfig.from_password(TEST_HOST, TEST_USER,
TEST_PASSWORD)
assert config.is_valid_with_password
assert not config.is_valid_with_token
def test_from_password(self):
config = DatabricksConfig.from_password(TEST_HOST, TEST_USER,
TEST_PASSWORD)
assert config.host == TEST_HOST
assert config.username == TEST_USER
assert config.password == TEST_PASSWORD
def get_config(self):
return DatabricksConfig.from_password("host",
"user",
"pass",
insecure=False)