CODEWITHSUNDEEP
pythonairflow
Eugene Brown
Eugene Brown

Reputation: 4362

Set dag_run.conf parameters in call to airflow test

Does anyone know if there is a way to set dag_run.conf parameters when running airflow test in the bash prompt?

For example, I've downloaded the example_trigger_target_dag from the official airflow repository and I'd like to test the run_this task. Usually I would do the following:

~/$ airflow test example_trigger_target_dag run_this '2018-01-01'

However running this produces the error:

--------------------------------------------------------------------------------
Starting attempt 1 of 1
--------------------------------------------------------------------------------

[2018-05-02 10:50:01,154] {models.py:1342} INFO - Executing <Task(PythonOperator): run_this> on 2018-01-01 00:00:00
[2018-05-02 10:50:01,262] {models.py:1417} ERROR - 'NoneType' object has no attribute 'conf'
Traceback (most recent call last):
  File "/usr/local/lib/python2.7/dist-packages/airflow/models.py", line 1374, in run
    result = task_copy.execute(context=context)
  File "/usr/local/lib/python2.7/dist-packages/airflow/operators/python_operator.py", line 80, in execute
    return_value = self.python_callable(*self.op_args, **self.op_kwargs)
  File "/home/annalect/uk_ds_airflow/dags/playpen/example_trigger_target_dag.py", line 56, in run_this_func
    print("Remotely received value of {} for key=message".format(kwargs['dag_run'].conf['message']))
AttributeError: 'NoneType' object has no attribute 'conf'

I've had a go at using the task_params argument however I either have the wrong syntax or it doesn't achieve what I'm after as it produces the same error as above:

~/$ airflow test --task_params '{"kwargs": {"dag_run": {"conf": {"message": "Hey world"}}}}' example_trigger_target_dag run_this '2018-01-01'

[2018-05-02 11:10:58,065] {models.py:1441} INFO - Marking task as FAILED.
[2018-05-02 11:10:58,070] {models.py:1462} ERROR - 'NoneType' object has no attribute 'conf'

So does anyone know how to test a task that depends on dag_run.conf values?

Thanks!

Upvotes: 13

Views: 18899

Answers (1)

Jacques Gaudin
Jacques Gaudin

Reputation: 16958

There is no --conf option for the airflow test command but you can work around this by passing parameters to the task's python_callable.

In the callable, if kwargs['test_mode'] is set, you can retrieve the parameters to build a dummy DagRun object like so:

from airflow.models import DagRun
...

def run_this_func(ds, **kwargs):
    if kwargs['test_mode']:
        kwargs['dag_run'] = DagRun(conf=kwargs['params'])

    print("Remotely received value of {} for key=message".format(kwargs['dag_run'].conf['message']))

To test example_trigger_target_dag, simply do:

airflow test example_trigger_target_dag test_trigger_dagrun "2018-01-01" -tp '{"message":"Hello world"}'

and you will get:

Remotely received value of Hello world for key=message

Now rather than putting test code in your tasks, you can write a decorator. Also since we are just using the conf attribute of DagRun, we may just as well use a SimpleNamespace. And finally, to avoid a potential key error whilst looking-up kwargs, we can use get with a default value.

from types import SimpleNamespace

def allow_conf_testing(func):
    def wrapper(*args, **kwargs):
        if kwargs.get('test_mode', False):
            kwargs['dag_run'] = SimpleNamespace(conf=kwargs.get('params', {}))
        func(*args, **kwargs)
    return wrapper

@allow_conf_testing
def run_this_func(ds, **kwargs):
    print("Remotely received value of {} for key=message".format(kwargs['dag_run'].conf['message']))

Upvotes: 16

Related Questions