def set_dbnd_unit_test_mode():
set_on(ENV_DBND__UNITTEST_MODE) # bypass to subprocess
get_dbnd_project_config().unit_test_mode = True
def disabled(self, value):
set_on(ENV_DBND__DISABLED)
self._disabled = value
ENV_DBND__NO_PLUGINS,
reset_dbnd_project_config,
)
from dbnd._core.plugin.dbnd_plugins import disable_airflow_plugin
from dbnd._core.utils.basics.environ_utils import set_on
from dbnd.testing.test_config_setter import add_test_configuration
from dbnd_test_scenarios.test_common.task.factories import FooConfig, TConfig
from targets import target
# we want to test only this module
# However, this import runs in separate space from test run -> this global will not be visible to your test
# get_dbnd_project_config().is_no_modules = True
# if enabled will pring much better info on tests
# os.environ["DBND__VERBOSE"] = "True"
set_on(ENV_DBND__NO_MODULES)
set_on(ENV_DBND__NO_PLUGINS)
set_on(ENV_DBND__DISABLE_PLUGGY_ENTRYPOINT_LOADING)
# DISABLE AIRFLOW, we don't test it in this module!
disable_airflow_plugin()
# all env changes should be active for current project config
reset_dbnd_project_config()
pytest_plugins = [
"dbnd.testing.pytest_dbnd_plugin",
"dbnd.testing.pytest_dbnd_markers_plugin",
"dbnd.testing.pytest_dbnd_home_plugin",
]
__all__ = ["dbnd"]
def set_context(self, ti):
"""
Airflow's log handler use this method to setup the context when running a TaskInstance(=ti).
We use this method to setup the dbnd context and communicate information to
the `<airflow_operator>_execute` task, that we create in `execute_tracking.py`.
"""
# we setting up only when we are not in our own orchestration dag
if ti.dag_id.startswith(AD_HOC_DAG_PREFIX):
return
if not is_dag_eligable_for_tracking(ti.dag_id):
return
if config.getboolean("mlflow_tracking", "databand_tracking"):
self.airflow_logger.warning(
"dbnd can't track mlflow and airflow together please disable dbnd config "
"`databand_tracking` in section `mlflow_tracking`")
return
# we are not tracking SubDagOperator
if ti.operator is None or ti.operator == SubDagOperator.__name__:
return
# Airflow is running with two process `run` and `--raw run`.
# But we want the handler to run only once (Idempotency)
# So we are using an environment variable to sync those two process
task_key = calc_task_key_from_af_ti(ti)
if os.environ.get(task_key, False):
# This key is already set which means we are in `--raw run`
return
else:
# We are in the outer `run`
self.task_env_key = task_key
# marking the environment with the current key for the
environ_utils.set_on(task_key)
from dbnd_airflow.tracking.dbnd_airflow_conf import (
set_dbnd_config_from_airflow_connections, )
# When we are in `--raw run`, in tracking, it runs the main airflow process
# for every task, which made some of the features to run twice,
# once when the `worker` process ran, and once when the `main` one ran,
# which made some of the features to run with different configurations.
# it still runs twice, but know with the same configurations.
set_dbnd_config_from_airflow_connections()
self.task_run_attempt_uid = get_task_run_attempt_uid_from_af_ti(ti)
# airflow calculation for the relevant log_file
log_relative_path = self.log_file_name_factory(ti, ti.try_number)
self.log_file = os.path.join(self.airflow_base_log_dir,
log_relative_path)
# make sure we are not polluting the airflow logs
get_dbnd_project_config().quiet_mode = True
# tracking msg
self.airflow_logger.info("Databand Tracking Started {version}".format(
version=dbnd.__version__))
# context with disabled logs
self.dbnd_context_manage = new_dbnd_context(
conf={"log": {
"disabled": True
}})
self.dbnd_context = self.dbnd_context_manage.__enter__()
)
from dbnd._core.utils.basics.environ_utils import set_on
from dbnd.testing.test_config_setter import add_test_configuration
from .mock_airflow_data_fetcher import MockDataFetcher
from .mock_airflow_tracking_service import MockServersConfigService, MockTrackingService
home = os.path.abspath(
os.path.normpath(os.path.join(os.path.dirname(__file__), "home")))
os.environ["DBND_HOME"] = home
os.environ["AIRFLOW_HOME"] = home
os.environ["AIRFLOW__CORE__UNIT_TEST_MODE"] = "True"
reset_dbnd_project_config()
# we don't need to load dbnd plugins/modules
set_on(ENV_DBND__NO_PLUGINS)
def pytest_configure(config):
add_test_configuration(__file__)
@fixture
def unittests_db():
return "fetch-unittests.db"
@fixture
def empty_db():
return "empty-unittests.db"