How to run an Airflow DAG that is defined in a test?

Viewed 499

I am trying to write a test case where I:

  • instantiate a collection of (Python)Operators (patching some of their dependencies with unittest.mock.patch)
  • arrange those Operators in a DAG
  • run that DAG
  • make assertions about the calls to various mocked downstreams

I see from here that running a DAG is not so simple as calling dag.run - I should instantiate a local_client and call trigger_dag on that. However, the resultant code constructs its own DagBag, and does not accept any parameter that allows me to pass in my manually-constructed DAG - so I cannot see how to run this DAG with local_client.

I see a couple of options here:

  • I could declare my testing DAG in a separate folder, as specified by DagModel.get_current(dag_id).fileloc, so that my DAG will get picked up by trigger_dag and so run - but this seems pretty indirect, and also I doubt that I'd be able to cleanly reference the injected mocks from my test code.
  • I could directly call api.common.experimental.trigger_dag._trigger_dag, which has a dag_bag argument. Both the experimental in the name, and the underscored-prefixed method name, suggest that this would be A Bad Idea.
0 Answers
Related