This project demonstrates how we can use docker-compose with pytest to run an integration test suite for Airflow DAGs. There is an article now in Medium detailing also manual tests with Jupyter notebook.

This is based on a heavily modified version of this useful article Airflow Integration Testing using Docker Compose

also I have thrown in the mix parts of the circle of ideas outlined here:

My Medium article

The list of changes follow

1. We use the official docker compose
2. We add the Mongo db with proper health checks and volumes
3. We add a jupyter notebook for experimenting with Mongo
4. Integration testing is done outside of the image
5. The sample dag also creates documents in the collection, currently commented out
6. Testcontainers execution, no need to run docker compose manually!!!
7. Replace custom triggering code with official airflow client calls in a more robust way.

We add an environment variable in docker compose according to Connections for the mongo hook

AIRFLOW_CONN_MONGO_STORE: 'mongodb://myTester:tester123@mongodb:27017/test?authSource=test'

which is referred by mongo_store. We also add a variable for the document collection

AIRFLOW_VAR_DOCUMENT_COLLECTION: 'mytest_collection'

called mytest_collection.

Create a virtual environment (tested with 3.11.3) and use requirements-dev.txt. Integration testing cannot be simpler than a

make integration_test

It starts a modified deployment, creates a test user, it runs the test, it deletes the user and kills the deployment. All automatic through testcontainers-python.

It is actually an automation of test_sample_dag.py (in the original article) as test_automatically_saple_dag.py.

Start the deployment

docker-compose up

If you are interested in making stateful, bring back the commented out volumes for postgres and mongo.

Now time to create our user (see docker file for the user). Connect to mongo

docker exec -it airflow-integration-testing-mongodb-1 mongosh -u root -p example

Use the following admin command in mongo to create the user, taken from here

use test
db.createUser(
  {
    user: "myTester",
    pwd:  passwordPrompt(),   // or cleartext password
    roles: [ { role: "readWrite", db: "test" },
             { role: "read", db: "reporting" } ]
  }
)

Set as password tester123. Exit and you can verify it works if you can connect with

docker exec -it airflow-integration-testing-mongodb-1 mongosh --port 27017 -u "myTester" --authenticationDatabase "test" -p

You can also use the notebook of jupyter. Run

docker logs airflow-integration-testing-datascience-notebook-1

to get the token/url. Just change port 8888 to 8891 and you are good to go.

Now from Airflow http://localhost:8080 activate you sample dag.

Read the corresponding operator. In summary every 2 minutes it checks if a document exists with {"source": 'some_source'} and marks it as processed. For debugging purposes there is a commented out section that actually creates a document.