如何修补/模拟导入？

Question

我正在为 airflow dag 编写测试并遇到模拟/修补 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

因为我只是在 function 或 class 之外设置变量，所以它在导入期间运行 Variable.get()。 这将给我一个 SQLAlchemy 错误，因为它正在尝试连接到数据库并获取变量。

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)]

有没有办法在导入之前修补/模拟 airflow.models.Variable？

Answer 1

您需要推迟导入文件，直到您可以将Variable值设置到测试数据库中。 startTestRun方法将是一个完美的地方。

Answer 2

我有一些解决方法，但没有明确的答案：

1. 您可以将此行ENVIRONMENT = Variable.get("environment")移动到 function 内部，而不是全局。 这样，它不会在导入时执行，您可以将此模拟添加到conftest.py ：

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

2. 您可以在测试 function 中导入模块。这样，将在调用导入之前设置模拟。

Answer 3

您可以在导入文件之前模拟导入

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()