使用 get_current_context 进行单元测试

Question

我想为在 Apache Airflow 中使用 get_current_context 的函数构建一个单元测试。该函数在多个任务中使用来创建一个文件名，用于从这些不同的任务中读取和写入文件。

这是该函数的示例：

def get_filename():
    from airflow.operators.python import get_current_context

    context = get_current_context()
    dag_id = context['dag'].__dict__['_dag_id']
    log_time = context['data_interval_start'].strftime("%Y-%m-%d_%H-%M-%S")
    log_file = f'/path/logs/{dag_id}/{log_time}.txt'
    return log_file

如何在单元测试中设置上下文以使该函数可执行？我什至不知道从哪里开始。

Answer 1

您需要模拟外部库函数调用。

from unittest import mock

from your.file.path import get_filename

with mock.patch('your.file.path.get_current_context') as mock_get_current_context:
    mock_get_current_context.return_value = <define the return value object>
    actual_filename = get_filename()
    expected_filename = "/path/logs/<some_id_which_was_used_in_mock_return_value>/<time_used_in_mock_return_value>.txt"
    assert actual_filename == expected_filename

要了解有关模拟和补丁的更多信息，请参阅https://docs.python.org/3/library/unittest.mock-examples.html