如何修補/模擬導入？

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