How to patch/mock import?

Question

I'm writing tests for airflow dag and running into issue mocking/patching the dag.

# dag.py
from airflow.models import Variable

ENVIRONMENT = Variable.get("environment")

# test_dag.py
import dag

class TestDAG(TestCase):
    def test_something(self):
        pass

Because I'm just setting variable outside of function or class, it runs Variable.get() during import. This will be give me a SQLAlchemy error cause it's trying to connect to a db and fetch variable.

sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) no such table: variable
[SQL: SELECT variable.val AS variable_val, variable.id AS variable_id, variable."key" AS variable_key, variable.is_encrypted AS variable_is_encrypted 
FROM variable 
WHERE variable."key" = ?
 LIMIT ? OFFSET ?]
[parameters: ('environment', 1, 0)]

Is there a way to patch/mock airflow.models.Variable before it's imported?

Accelerometer

I ended up patching it before import which isn't great but I haven't found other ways. This produces pep8 and linting errors because imports aren't grouped at top anymore. I'll leave the question open in case anyone has suggestions – Accelerometer

Could you show that in code by any chance, please? – Globigerina

Please show us your workaround. I have the same issue and is pretty annoying. – Testudinal

I'm having the same problem, someone resolved this? – Pushup

Answer 1

You'll need to defer importing the file until you can set a Variable value into the test database. A startTestRun method would be the perfect place.

Mussman

I don't want to connect to a test database, I want to mock it. Could you provide an example? I can't get it to work – Accelerometer

Often developers will use an in-memory SQLite database that is built/destroyed for each test as needed. – Mussman

Answer 2

I have some workarounds but not a definite answer:

1. You can move this line ENVIRONMENT = Variable.get("environment") into inside a function, instead of global. This way, it will not be executed when imported and you can add this mock to conftest.py:

@pytest.fixture(autouse=True)
def mock_airflow_variables(mocker):
        mocker.patch.object(target=Variable, attribute="get", return_value="test")

2. You can import the module inside the test function. This way, the mock will be set before calling import.

Pushup

Answer 3

You can mock the import before importing the file

import unittest
from unittest.mock import MagicMock

import sys


class TestDAG(unittest.TestCase):

    def test_something(self):
        sys.modules['airflow.models'] = MagicMock()
    
        # This returns the MagicMock instance
        from airflow.models import Variable 

        # Set the return_value of the .get() call
        Variable.get.return_value = "TEST" 

        # import the dag after
        import dag

        Variable.get.assert_called_once()
        assert dag.ENVIRONMENT == "TEST"


if __name__ == '__main__':
    unittest.main()

Neidaneidhardt

Answer 4

Here is a working example:

from airflow.models import abstractoperator, baseoperator, dag, xcom_arg, skipmixin

sys.modules["airflow.models"] = MagicMock()
sys.modules["airflow.models"].abstractoperator.return_value = abstractoperator
sys.modules["airflow.models"].baseoperator.return_value = baseoperator
sys.modules["airflow.models"].dag.return_value = dag
sys.modules["airflow.models"].xcom_arg.return_value = xcom_arg
sys.modules["airflow.models"].skipmixin.return_value = skipmixin

To avoid import failures I had to set the return_value back to the original imports.

Variable for example is now mocked.

apache-airflow==2.7.3

Sociometry