[Date Prev][Date Next][Thread Prev][Thread Next][Date Index][Thread Index]

Mocking airflow (similar to moto for AWS)

I think it would be great to have a way to mock airflow for unit tests. The way I approached this was to create a context manager that creates a temporary directory, sets the AIRFLOW_HOME environment variable to this directory (only within the scope of the context manager) and then renders an airflow.cfg to that location. This creates an SQLite just for the test so you can add variables and connections needed for the test without affecting the real Airflow installation.

The first thing I realized is that this didn't work if the imports were outside the context manager, since airflow.configuration and airflow.settings perform all the initialization when they are imported, so the AIRFLOW_HOME variable is already set to the real installation before getting inside the context manager.

The workaround for this was to reload those modules and this works for the tests I have written. However, when I tried to use it for something more complex (I have a plugin that I'm importing) I noticed that inside the operator in this plugin, AIRFLOW_HOME is still set to the real installation, not the temporary one for the test. I thought this must be related to the imports but I haven't been able to figure out a way to fix the issue. I tried patching some methods but I must have been missing something because the database initialization failed.

Does anyone have an idea on the best way to mock/patch airflow so that EVERYTHING that is executed inside the context manager uses the temporary installation?

PS: This is my current attempt which works for the tests I defined but not for external plugins: 

For an example on how it works: https://github.com/biellls/airflow_testing/blob/master/tests/mock_airflow_test.py