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)