Unit testing with get_current_context

user_17359519 0
Viewed 375

I want to build a unit test for a function which uses get_current_context in Apache Airflow. The function is used within multiple tasks to create a filename used to read and write to the file from these different tasks.

Here is an example of the function:

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

How do I set the context in the unit test so that the function is executable? I don't even know where to begin.

airflow unit-testing python
1 Answers

You need to mock external library function calls.

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

To learn more about mock and patch, refer to https://docs.python.org/3/library/unittest.mock-examples.html

user_46595300
Related