Writing integration tests for DAGs

[FredericoCoelhoNunes]

Hi everyone,

so I am currently trying to write tests for Airflow DAGs using the unittest package and running into some conceptual issues.

My code is basically the following:

I have a function that returns my DAG. This way I can instantiate it once (the production DAG) and ideally, instantiate a temporary test version of it in the test suite. The reasoning behind this is that I don't want to add runs to the production DAG every time I run the test suite.

def my_dag(dag_id, input_data, output_location):
    dag = DAG(dag_id=dag_id)
    task_1 = ...
    task_2 = ...
    task_1 >> task_2
    return dag

This DAG is instantiated inside my DAGS_FOLDER (all good so far).

Then inside my test file I am doing the following:

from airflow.api.client.local_client import Client
from airflow.models.dagbag import DagBag

...

def test_my_dag():

        test_dag = my_dag(
            dag_id = 'test-dag'
            input_data = my_test_input_data,
            output_location = test_output_location
        )

        <something goes in here...>
        
        c.trigger_dag(
            dag_id="test-dag",
            run_id="test-run-id",
            conf={}
        )

I have tried several things for the section, but no luck so far. I keep getting airflow.exceptions.DagNotFound: Dag id test-dag not found in DagModel . I've tried instantiating the DagBag and adding the dag to it (using the bag_dag method or by manually setting dagbag.dags[test_dag.dag_id] = test_dag), but no luck.

What I wanted was a way to create an ephemeral DAG whose only purpose is to be instantiated and run during the test! This seemed like a reasonable software engineering approach.

Has anyone faced this problem in the past? How did you solve it? Any suggestions are very welcome!

Cheers

[potiuk]

You can take a look how we implemented System Tests in Apache Airflow : https://github.com/apache/airflow/blob/master/TESTING.rst#airflow-system-tests

System tests are basically executing an 'example_dag' and wait until it completes. We have a lot of such example DAGs in Airflow and there is an effort ongoing to automate execution of those #10835

This is very much your case under the hood - but done using some of the Airflow classes - loading DAG is done via DagBag and then it is executed using a Local (i think? Maybe it was the Debug one) Executor. All that is packed in SystemTest class that does extra initialization and allows to run it as Pytest test case.

[FredericoCoelhoNunes]

Hi @potiuk

thanks for the help! I checked your code and realized I completely missed the dag.run(...) method.

Right now I'm trying to trigger it this way, but it apparently is trying to backfill the DAG even though I set catchup to False (I'm certain I'm missing something).

The only difference from my code to yours is that I am instantiating the DAG in the code and manually adding it to the DagBag, whereas you have the example_dags in a folder and the DagBag picks them up as soon as it is instantiated. In my case, I have to create it as DagBag("fake folder") which is a bit ugly, but gives me an empty DagBag to play around with.

Anyway, thanks a lot for your lightning quick feedback, it was very helpful indeed.

Cheers

[potiuk]

I think the trick will be to set those in your DAG:

    schedule_interval=None, 
    start_date=days_ago(1),

That should help with the back-fill problem.

[FredericoCoelhoNunes]

Thanks @potiuk !