def test_exec_failure(self, db_mock_class):
"""
Test the execute function in case where the run failed.
"""
run = {
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
}
op = DatabricksSubmitRunOperator(task_id=TASK_ID, json=run)
db_mock = db_mock_class.return_value
db_mock.submit_run.return_value = 1
db_mock.get_run_state.return_value = RunState('TERMINATED', 'FAILED',
'')
with self.assertRaises(AirflowException):
op.execute(None)
expected = op._deep_string_coerce({
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
'run_name': TASK_ID,
})
db_mock_class.assert_called_once_with(
DEFAULT_CONN_ID, retry_limit=op.databricks_retry_limit)
db_mock.submit_run.assert_called_once_with(expected)
db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
db_mock.get_run_state.assert_called_once_with(RUN_ID)
self.assertEquals(RUN_ID, op.run_id)
def test_exec_failure(self, db_mock_class):
"""
Test the execute function in case where the run failed.
"""
run = {
'notebook_params': NOTEBOOK_PARAMS,
'notebook_task': NOTEBOOK_TASK,
'jar_params': JAR_PARAMS
}
op = DatabricksRunNowOperator(task_id=TASK_ID, job_id=JOB_ID, json=run)
db_mock = db_mock_class.return_value
db_mock.run_now.return_value = 1
db_mock.get_run_state.return_value = RunState('TERMINATED', 'FAILED',
'')
with self.assertRaises(AirflowException):
op.execute(None)
expected = databricks_operator._deep_string_coerce({
'notebook_params': NOTEBOOK_PARAMS,
'notebook_task': NOTEBOOK_TASK,
'jar_params': JAR_PARAMS,
'job_id': JOB_ID
})
db_mock_class.assert_called_once_with(
DEFAULT_CONN_ID,
retry_limit=op.databricks_retry_limit,
retry_delay=op.databricks_retry_delay)
db_mock.run_now.assert_called_once_with(expected)
db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
db_mock.get_run_state.assert_called_once_with(RUN_ID)
self.assertEqual(RUN_ID, op.run_id)
def test_exec_success(self, db_mock_class):
"""
Test the execute function in case where the run is successful.
"""
run = {
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
}
op = DatabricksSubmitRunOperator(task_id=TASK_ID, json=run)
db_mock = db_mock_class.return_value
db_mock.submit_run.return_value = 1
db_mock.get_run_state.return_value = RunState('TERMINATED', 'SUCCESS',
'')
op.execute(None)
expected = databricks_operator._deep_string_coerce({
'new_cluster': NEW_CLUSTER,
'notebook_task': NOTEBOOK_TASK,
'run_name': TASK_ID
})
db_mock_class.assert_called_once_with(
DEFAULT_CONN_ID,
retry_limit=op.databricks_retry_limit,
retry_delay=op.databricks_retry_delay)
db_mock.submit_run.assert_called_once_with(expected)
db_mock.get_run_page_url.assert_called_once_with(RUN_ID)
db_mock.get_run_state.assert_called_once_with(RUN_ID)
self.assertEqual(RUN_ID, op.run_id)
def test_simple_python(self, d_hook):
run_states = [
RunState("RUNNING", "gla", "bla"),
RunState("TERMINATED", "SUCCESS", "yey"),
]
d_hook_instance = d_hook.return_value
d_hook_instance.submit_run = my_submit_run
d_hook_instance.get_run_state.side_effect = run_states
d_hook_instance.get_run_page_url.return_value = (
"https://databand.ai-test-databricks")
WordCountPySparkTask(
text=__file__,
spark_engine="databricks",
task_version="now",
override=test_config,
).dbnd_run()
def test_get_run_state(self, mock_requests):
mock_requests.get.return_value.json.return_value = GET_RUN_RESPONSE
run_state = self.hook.get_run_state(RUN_ID)
self.assertEquals(
run_state, RunState(LIFE_CYCLE_STATE, RESULT_STATE, STATE_MESSAGE))
mock_requests.get.assert_called_once_with(
get_run_endpoint(HOST),
json={'run_id': RUN_ID},
auth=(LOGIN, PASSWORD),
headers=USER_AGENT_HEADER,
timeout=self.hook.timeout_seconds)
def test_is_successful(self):
run_state = RunState('TERMINATED', 'SUCCESS', '')
self.assertTrue(run_state.is_successful)
def test_is_terminal_with_nonexistent_life_cycle_state(self):
run_state = RunState('blah', '', '')
with self.assertRaises(AirflowException):
run_state.is_terminal
def test_is_terminal_false(self):
non_terminal_states = ['PENDING', 'RUNNING', 'TERMINATING']
for state in non_terminal_states:
run_state = RunState(state, '', '')
self.assertFalse(run_state.is_terminal)
def test_is_terminal_true(self):
terminal_states = ['TERMINATED', 'SKIPPED', 'INTERNAL_ERROR']
for state in terminal_states:
run_state = RunState(state, '', '')
self.assertTrue(run_state.is_terminal)
