def test_cancel_cli():
with mock.patch('databricks_cli.runs.cli.cancel_run') as cancel_run_mock:
with mock.patch('databricks_cli.runs.cli.click.echo') as echo_mock:
cancel_run_mock.return_value = {}
get_callback(cli.cancel_cli)(1)
assert cancel_run_mock.call_args[0][0] == 1
assert echo_mock.call_args[0][0] == pretty_format({})
def test_reset_cli_json():
with mock.patch('databricks_cli.jobs.cli.reset_job') as reset_job_mock:
get_callback(cli.reset_cli)(None, RESET_JSON, 1)
assert reset_job_mock.call_args[0][0] == {
'job_id': 1,
'new_settings': json.loads(RESET_JSON)
}
def test_create_cli_json():
with mock.patch('databricks_cli.jobs.cli.create_job') as create_job_mock:
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
create_job_mock.return_value = CREATE_RETURN
get_callback(cli.create_cli)(None, CREATE_JSON)
assert create_job_mock.call_args[0][0] == json.loads(CREATE_JSON)
assert echo_mock.call_args[0][0] == pretty_format(CREATE_RETURN)
def test_submit_cli_json():
with mock.patch('databricks_cli.runs.cli.submit_run') as submit_run_mock:
with mock.patch('databricks_cli.runs.cli.click.echo') as echo_mock:
submit_run_mock.return_value = SUBMIT_RETURN
get_callback(cli.submit_cli)(None, SUBMIT_JSON)
assert submit_run_mock.call_args[0][0] == json.loads(SUBMIT_JSON)
assert echo_mock.call_args[0][0] == pretty_format(SUBMIT_RETURN)
def test_list_clusters_output_json():
with mock.patch(
'databricks_cli.clusters.cli.list_clusters') as list_clusters_mock:
with mock.patch('databricks_cli.clusters.cli.click.echo') as echo_mock:
list_clusters_mock.return_value = LIST_RETURN
get_callback(cli.list_cli)('json')
assert echo_mock.call_args[0][0] == pretty_format(LIST_RETURN)
def test_list_runs():
with mock.patch('databricks_cli.runs.cli.list_runs') as list_runs_mock:
with mock.patch('databricks_cli.runs.cli.click.echo') as echo_mock:
list_runs_mock.return_value = LIST_RETURN
get_callback(cli.list_cli)(None, None, None, None, None, None)
rows = [(1, 'name', 'RUNNING', 'n/a', RUN_PAGE_URL)]
assert echo_mock.call_args[0][0] == tabulate(rows,
tablefmt='plain')
def test_list_jobs():
with mock.patch(
'databricks_cli.clusters.cli.list_clusters') as list_clusters_mock:
with mock.patch('databricks_cli.clusters.cli.click.echo') as echo_mock:
list_clusters_mock.return_value = LIST_RETURN
get_callback(cli.list_cli)(None)
assert echo_mock.call_args[0][0] == \
tabulate([('test_id', 'test_name', 'PENDING')], tablefmt='plain')
def test_list_jobs():
with mock.patch('databricks_cli.jobs.cli.list_jobs') as list_jobs_mock:
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
list_jobs_mock.return_value = LIST_RETURN
get_callback(cli.list_cli)(None)
# Output should be sorted here.
rows = [(2, 'a'), (1, 'b'), (30, 'C')]
assert echo_mock.call_args[0][0] == \
tabulate(rows, tablefmt='plain', disable_numparse=True)
def test_create_cli_json_file(tmpdir):
path = tmpdir.join('job.json').strpath
with open(path, 'w') as f:
f.write(CREATE_JSON)
with mock.patch('databricks_cli.jobs.cli.create_job') as create_job_mock:
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
create_job_mock.return_value = CREATE_RETURN
get_callback(cli.create_cli)(path, None)
assert create_job_mock.call_args[0][0] == json.loads(CREATE_JSON)
assert echo_mock.call_args[0][0] == pretty_format(CREATE_RETURN)
def test_run_now_no_params():
with mock.patch('databricks_cli.jobs.cli.run_now') as run_now_mock:
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
run_now_mock.return_value = RUN_NOW_RETURN
get_callback(cli.run_now_cli)(1, None, None, None, None)
assert run_now_mock.call_args[0][0] == 1
assert run_now_mock.call_args[0][1] is None
assert run_now_mock.call_args[0][2] is None
assert run_now_mock.call_args[0][3] is None
assert run_now_mock.call_args[0][4] is None
assert echo_mock.call_args[0][0] == pretty_format(RUN_NOW_RETURN)
def test_run_now_with_params():
with mock.patch('databricks_cli.jobs.cli.run_now') as run_now_mock:
with mock.patch('databricks_cli.jobs.cli.click.echo') as echo_mock:
run_now_mock.return_value = RUN_NOW_RETURN
get_callback(cli.run_now_cli)(1, JAR_PARAMS, NOTEBOOK_PARAMS, PYTHON_PARAMS,
SPARK_SUBMIT_PARAMS)
assert run_now_mock.call_args[0][0] == 1
assert run_now_mock.call_args[0][1] == json.loads(JAR_PARAMS)
assert run_now_mock.call_args[0][2] == json.loads(NOTEBOOK_PARAMS)
assert run_now_mock.call_args[0][3] == json.loads(PYTHON_PARAMS)
assert run_now_mock.call_args[0][4] == json.loads(SPARK_SUBMIT_PARAMS)
assert echo_mock.call_args[0][0] == pretty_format(RUN_NOW_RETURN)
def test_export_workspace_cli(tmpdir):
path = tmpdir.strpath
with mock.patch(
'databricks_cli.workspace.cli.get_status') as get_status_mock:
with mock.patch('databricks_cli.workspace.cli.export_workspace'
) as export_workspace_mock:
get_status_mock.return_value = WorkspaceFileInfo(
'/notebook-name', NOTEBOOK, WorkspaceLanguage.SCALA)
get_callback(cli.export_workspace_cli)('/notebook-name', path,
WorkspaceFormat.SOURCE,
False)
assert export_workspace_mock.call_args[0][1] == os.path.join(
path, 'notebook-name.scala')
def test_get_cli():
with mock.patch(
'databricks_cli.clusters.cli.get_cluster') as get_cluster_mock:
get_cluster_mock.return_value = '{}'
get_callback(cli.get_cli)(CLUSTER_ID)
assert get_cluster_mock.call_args[0][0] == CLUSTER_ID
def test_delete_cli():
with mock.patch('databricks_cli.clusters.cli.delete_cluster'
) as delete_cluster_mock:
get_callback(cli.delete_cli)(CLUSTER_ID)
assert delete_cluster_mock.call_args[0][0] == CLUSTER_ID
def test_restart_cli():
with mock.patch('databricks_cli.clusters.cli.restart_cluster'
) as restart_cluster_mock:
get_callback(cli.restart_cli)(CLUSTER_ID)
assert restart_cluster_mock.call_args[0][0] == CLUSTER_ID