def test_06_both_0(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(with_log=True)
result = get_new_dag_runs(0, 0, [])
self.validate_result(result, 1, 1, 1, expected_max_log_ids=[1])
def test_02_get_task_instances(self):
from dbnd_airflow.export_plugin.api_functions import get_dag_runs_states_data
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, task_instances_per_run=3)
result = get_dag_runs_states_data([1, 2, 3])
self.validate_result(result, 9)
def test_03_running(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, state="running", with_log=False)
result = get_new_dag_runs(None, None, [])
self.validate_result(result, 3, 3, None)
def test_05_log_id_none(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=True)
result = get_new_dag_runs(1, None, [])
self.validate_result(result, 2, 3, 3)
def test_07_big_run_id(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=False)
result = get_new_dag_runs(3, 0, [])
self.validate_result(result, 0, 3, None)
assert result
def test_10_running_paused(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import (
insert_dag_runs,
set_dag_is_paused,
)
insert_dag_runs(dag_runs_count=1, state="running", with_log=True)
set_dag_is_paused(is_paused=True)
result = get_new_dag_runs(1, 1, [])
self.validate_result(result, 0, 1, 1, True)
def test_09_paused(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import (
insert_dag_runs,
set_dag_is_paused,
)
insert_dag_runs(dag_runs_count=1, with_log=True)
set_dag_is_paused(is_paused=True)
result = get_new_dag_runs(0, 0, [])
self.validate_result(result, 1, 1, 1, True, expected_max_log_ids=[1])
def test_12_dag_ids(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=True)
insert_dag_runs(dag_id="plugin_other_dag",
dag_runs_count=3,
with_log=True)
result = get_new_dag_runs(0, 0, [], ["plugin_other_dag"])
self.validate_result(result, 3, 6, 6, expected_max_log_ids=[4, 5, 6])
def test_11_extra_dag_runs(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=True)
result = get_new_dag_runs(2, 2, [1, 2])
self.validate_result(result,
3,
3,
3,
expected_max_log_ids=[None, None, 3])
def test_02_get_task_instances(self):
from dbnd_airflow.export_plugin.api_functions import get_full_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, task_instances_per_run=3)
dag_run_ids = [1, 2, 3]
from dbnd_airflow.export_plugin.smart_dagbag import DbndDagLoader
dbnd_dag_loader = DbndDagLoader()
dbnd_dag_loader.load_dags_for_runs(dag_run_ids)
result = get_full_dag_runs(dag_run_ids, True, dbnd_dag_loader)
self.validate_result(result, 1, 3, 9)
def test_03_sync_no_sources(self):
from dbnd_airflow.export_plugin.api_functions import get_full_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, task_instances_per_run=3)
dag_run_ids = [1, 2, 3]
from dbnd_airflow.export_plugin.smart_dagbag import DbndDagLoader
dbnd_dag_loader = DbndDagLoader()
dbnd_dag_loader.load_dags_for_runs(dag_run_ids)
result = get_full_dag_runs(dag_run_ids, False, dag_loader=dbnd_dag_loader)
self.validate_result(result, 1, 3, 9)
for dag in result.dags:
assert not dag.source_code
for task in dag.tasks:
assert not task.task_source_code
assert not task.task_module_code
def test_13_fetch_in_chunks(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from dbnd_airflow.export_plugin.queries import (
MAX_PARAMETERS_INSIDE_IN_CLAUSE,
_find_dag_runs_by_list_in_chunks,
)
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
with mock.patch(
"dbnd_airflow.export_plugin.queries._find_dag_runs_by_list_in_chunks",
wraps=lambda a, b: _find_dag_runs_by_list_in_chunks(a, b),
) as m:
insert_dag_runs(dag_runs_count=MAX_PARAMETERS_INSIDE_IN_CLAUSE - 1,
with_log=True)
get_new_dag_runs(0, 0, [], [])
assert m.call_count == 0
insert_dag_runs(dag_runs_count=1, with_log=True)
get_new_dag_runs(0, 0, [], [])
assert m.call_count == 1
def test_dag_bag(self):
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
for i in range(1, 7):
insert_dag_runs(
dag_runs_count=1,
task_instances_per_run=3,
dag_id=f"plugin_test_dag_{i}",
)
from dbnd_airflow.export_plugin.smart_dagbag import DbndDagLoader
dbnd_dag_loader = DbndDagLoader()
dbnd_dag_loader.load_dags_for_runs([1, 2])
actual_ids = {d.dag_id for d in dbnd_dag_loader.get_dags().values()}
assert actual_ids == {"plugin_test_dag_1", "plugin_test_dag_2"}
dbnd_dag_loader.load_dags_for_runs([3, 4])
actual_ids = {d.dag_id for d in dbnd_dag_loader.get_dags().values()}
assert actual_ids == {
"plugin_test_dag_1",
"plugin_test_dag_2",
"plugin_test_dag_3",
"plugin_test_dag_4",
}
dbnd_dag_loader.load_dags_for_runs([5, 6])
actual_ids = {d.dag_id for d in dbnd_dag_loader.get_dags().values()}
assert actual_ids == {
"plugin_test_dag_1",
"plugin_test_dag_2",
"plugin_test_dag_3",
"plugin_test_dag_4",
"plugin_test_dag_5",
"plugin_test_dag_6",
}
