CODEWITHSUNDEEP
pythonunit-testingairflow
intervort
intervort

Reputation: 153

Unit testing with get_current_context

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.

Upvotes: 5

Views: 3107

Answers (1)

Rahul Kumar
Rahul Kumar

Reputation: 2345

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

Upvotes: 2

Related Questions