def test_tracking_pass_through_default_tracking( pandas_data_frame_on_disk, mock_channel_tracker ): df, df_file = pandas_data_frame_on_disk # we'll pass string instead of defined expected DataFrame and it should work some_date = utcnow().isoformat() task_result = task_pass_through_default( str(df_file), some_date, expect_pass_through=True ) assert task_result == str(df_file) # this should happen on process exit in normal circumstances dbnd_tracking_stop() _check_tracking_calls( mock_channel_tracker, { "init_run": 1, "add_task_runs": 1, "log_metrics": 1, "log_targets": 1, "save_task_run_log": 2, "set_run_state": 1, "update_task_run_attempts": 4, # DAG start(with task start), task started running, task finished, dag stop }, ) _assert_tracked_params( mock_channel_tracker, task_pass_through_default, data=str(df_file), dt=some_date, expect_pass_through=True, )
def set_tracking_context(): try: reset_dbnd_project_config() get_dbnd_project_config()._dbnd_tracking = True yield finally: dbnd_tracking_stop() reset_dbnd_project_config()
def test_tracking_pass_through_default_airflow( pandas_data_frame_on_disk, mock_channel_tracker ): df, df_file = pandas_data_frame_on_disk # we'll pass string instead of defined expected DataFrame and it should work from targets.values import DateTimeValueType some_date = DateTimeValueType().to_str(utcnow()) task_result = task_pass_through_default( str(df_file), some_date, expect_pass_through=True ) assert task_result == str(df_file) _check_tracking_calls( mock_channel_tracker, { "init_run": 1, "add_task_runs": 1, # real task only "update_task_run_attempts": 3, # DAG start(with task start), task started running, task finished "log_metrics": 1, "log_targets": 1, "save_task_run_log": 1, }, ) _assert_tracked_params( mock_channel_tracker, task_pass_through_default, data=str(df_file), dt=some_date, expect_pass_through=True, ) # this should happen on process exit in normal circumstances dbnd_tracking_stop() _check_tracking_calls( mock_channel_tracker, { "init_run": 1, "add_task_runs": 1, "update_task_run_attempts": 4, # as above + airflow root stop "log_metrics": 1, "log_targets": 1, "save_task_run_log": 2, # as above + airflow root log }, )
def set_airflow_context(): with patch( "dbnd._core.tracking.airflow_dag_inplace_tracking.try_get_airflow_context" ) as m: try: reset_dbnd_project_config() m.return_value = AirflowTaskContext( dag_id="test_dag", task_id="test_task", execution_date=utcnow().isoformat(), ) yield finally: # ensure dbnd_run_stop() is called (normally should happen on exit() ) dbnd_tracking_stop() reset_dbnd_project_config()
def test_tracking_keyboard_interrupt(mock_channel_tracker): # we'll pass string instead of defined expected DataFrame and it should work try: task_pass_through_keyboard_interrupt() except BaseException: _check_tracking_calls( mock_channel_tracker, { "init_run": 1, "add_task_runs": 1, "update_task_run_attempts": 3, "save_task_run_log": 1, }, ) # this should happen on process exit in normal circumstances dbnd_tracking_stop() _check_tracking_calls( mock_channel_tracker, { "init_run": 1, "add_task_runs": 1, "update_task_run_attempts": 4, "save_task_run_log": 2, }, ) update_task_run_attempts_chain = [ call.kwargs["task_run_attempt_updates"][0].state for call in mock_channel_tracker.call_args_list if call.args[0].__name__ == "update_task_run_attempts" ] assert [ TaskRunState.RUNNING, # DAG TaskRunState.RUNNING, # root_task TaskRunState.FAILED, # task TaskRunState.UPSTREAM_FAILED, # DAG ] == update_task_run_attempts_chain
def test_tracking_pass_through_nested_default( pandas_data_frame_on_disk, mock_channel_tracker ): df, df_file = pandas_data_frame_on_disk # we'll pass string instead of defined expected DataFrame and it should work task_result = task_pass_through_nested_default( str(df_file), utcnow().isoformat(), expect_pass_through=True ) assert task_result == str(df_file) + str(df_file) _check_tracking_calls( mock_channel_tracker, { "init_run": 1, "add_task_runs": 2, "update_task_run_attempts": 5, "log_metrics": 1, "log_targets": 2, "save_task_run_log": 2, }, ) # this should happen on process exit in normal circumstances dbnd_tracking_stop() _check_tracking_calls( mock_channel_tracker, { "init_run": 1, "add_task_runs": 2, "update_task_run_attempts": 6, "log_metrics": 1, "log_targets": 2, "save_task_run_log": 3, }, )
def test_sanity_with_airflow(self, with_airflow_tracking_env): fake_task_inside_dag() dbnd_tracking_stop()
@task def f1(): print("Running f1 function") sum = 0 for i in range(1, 4): sum += f2(i) assert sum == 14 if __name__ == "__main__": if USE_DBND_START in sys.argv: dbnd_tracking_start(conf={"log": {"disabled": False}}) dbnd_tracking_start(conf={"log": {"disabled": False}}) f1() print("Done") if FAIL_MAIN in sys.argv: raise Exception("main bummer!") if RUN_TASK_DECORATED_METHOD: MyClass().do_something() if USE_DBND_STOP in sys.argv: dbnd_tracking_stop() dbnd_tracking_stop()