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()