Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 15 additions & 9 deletions providers/sftp/src/airflow/providers/sftp/sensors/sftp.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,18 +95,24 @@ def poke(self, context: Context) -> PokeReturnValue | bool:
else:
return False
else:
actual_files_to_check = [self.path]

for actual_file_to_check in actual_files_to_check:
try:
mod_time = self.hook.get_mod_time(actual_file_to_check)
self.log.info("Found File %s last modified: %s", actual_file_to_check, mod_time)
self.hook.isfile(self.path)
actual_files_to_check = [self.path]
except OSError as e:
if e.errno != SFTP_NO_SUCH_FILE:
raise AirflowException from e
continue
actual_files_to_check = []

if self.newer_than:
Comment thread
kkulczak marked this conversation as resolved.
for actual_file_to_check in actual_files_to_check:
try:
mod_time = self.hook.get_mod_time(actual_file_to_check)
self.log.info("Found File %s last modified: %s", actual_file_to_check, mod_time)
except OSError as e:
if e.errno != SFTP_NO_SUCH_FILE:
raise AirflowException from e
continue

if self.newer_than:
if isinstance(self.newer_than, str):
self.newer_than = parse(self.newer_than)
_mod_time = convert_to_utc(datetime.strptime(mod_time, "%Y%m%d%H%M%S"))
Expand All @@ -126,8 +132,8 @@ def poke(self, context: Context) -> PokeReturnValue | bool:
str(_mod_time),
str(_newer_than),
)
else:
files_found.append(actual_file_to_check)
else:
files_found = actual_files_to_check

if not len(files_found):
return False
Expand Down
54 changes: 38 additions & 16 deletions providers/sftp/tests/unit/sftp/sensors/test_sftp.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

from datetime import datetime, timezone as stdlib_timezone
from unittest import mock
from unittest.mock import Mock, call, patch
from unittest.mock import Mock, patch

import pytest
from paramiko.sftp import SFTP_FAILURE, SFTP_NO_SUCH_FILE
Expand All @@ -36,27 +36,27 @@
class TestSFTPSensor:
@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_present(self, sftp_hook_mock):
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sftp_hook_mock.return_value.isfile.return_value = True
sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/1970-01-01.txt")
context = {"ds": "1970-01-01"}
output = sftp_sensor.poke(context)
sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.isfile.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert output

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_absent(self, sftp_hook_mock):
sftp_hook_mock.return_value.get_mod_time.side_effect = OSError(SFTP_NO_SUCH_FILE, "File missing")
sftp_hook_mock.return_value.isfile.side_effect = OSError(SFTP_NO_SUCH_FILE, "File missing")
sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/1970-01-01.txt")
context = {"ds": "1970-01-01"}
output = sftp_sensor.poke(context)
sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.isfile.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert not output

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_sftp_failure(self, sftp_hook_mock):
sftp_hook_mock.return_value.get_mod_time.side_effect = OSError(SFTP_FAILURE, "SFTP failure")
sftp_hook_mock.return_value.isfile.side_effect = OSError(SFTP_FAILURE, "SFTP failure")

sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/1970-01-01.txt")
context = {"ds": "1970-01-01"}
Expand All @@ -69,6 +69,7 @@ def test_hook_not_created_during_init(self):

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_new_enough(self, sftp_hook_mock):
sftp_hook_mock.return_value.isfile.return_value = True
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
tz = timezone("America/Toronto")
sftp_sensor = SFTPSensor(
Expand All @@ -78,12 +79,14 @@ def test_file_new_enough(self, sftp_hook_mock):
)
context = {"ds": "1970-01-00"}
output = sftp_sensor.poke(context)
sftp_hook_mock.return_value.isfile.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert output

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_not_new_enough(self, sftp_hook_mock):
sftp_hook_mock.return_value.isfile.return_value = True
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
tz = timezone("Europe/Paris")
sftp_sensor = SFTPSensor(
Expand All @@ -93,6 +96,7 @@ def test_file_not_new_enough(self, sftp_hook_mock):
)
context = {"ds": "1970-01-00"}
output = sftp_sensor.poke(context)
sftp_hook_mock.return_value.isfile.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert not output
Expand All @@ -113,30 +117,29 @@ def test_file_not_new_enough(self, sftp_hook_mock):
)
@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_multiple_datetime_format_in_newer_than(self, sftp_hook_mock, newer_than):
sftp_hook_mock.return_value.isfile.return_value = True
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sftp_sensor = SFTPSensor(
task_id="unit_test", path="/path/to/file/1970-01-01.txt", newer_than=newer_than
)
context = {"ds": "1970-01-00"}
output = sftp_sensor.poke(context)
sftp_hook_mock.return_value.isfile.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert not output

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_present_with_pattern(self, sftp_hook_mock):
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sftp_hook_mock.return_value.get_files_by_pattern.return_value = ["text_file.txt"]
sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/", file_pattern="*.txt")
context = {"ds": "1970-01-01"}
output = sftp_sensor.poke(context)
sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/text_file.txt")
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert output

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_not_present_with_pattern(self, sftp_hook_mock):
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sftp_hook_mock.return_value.get_files_by_pattern.return_value = []
sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/", file_pattern="*.txt")
context = {"ds": "1970-01-01"}
Expand All @@ -146,18 +149,14 @@ def test_file_not_present_with_pattern(self, sftp_hook_mock):

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_multiple_files_present_with_pattern(self, sftp_hook_mock):
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sftp_hook_mock.return_value.get_files_by_pattern.return_value = [
"text_file.txt",
"another_text_file.txt",
]
sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/", file_pattern="*.txt")
context = {"ds": "1970-01-01"}
output = sftp_sensor.poke(context)
get_mod_time = sftp_hook_mock.return_value.get_mod_time
expected_calls = [call("/path/to/file/text_file.txt"), call("/path/to/file/another_text_file.txt")]
sftp_hook_mock.return_value.close_conn.assert_not_called()
assert get_mod_time.mock_calls == expected_calls
assert output

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
Expand Down Expand Up @@ -227,7 +226,7 @@ def test_multiple_old_files_present_with_pattern_and_newer_than(self, sftp_hook_
)
@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_path_present_with_callback(self, sftp_hook_mock, op_args, op_kwargs):
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sftp_hook_mock.return_value.isfile.return_value = True
sample_callable = Mock()
sample_callable.return_value = ["sample_return"]
sftp_sensor = SFTPSensor(
Expand All @@ -240,7 +239,7 @@ def test_file_path_present_with_callback(self, sftp_hook_mock, op_args, op_kwarg
context = {"ds": "1970-01-01"}
output = sftp_sensor.poke(context)

sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.isfile.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.close_conn.assert_not_called()
sample_callable.assert_called_once_with(*op_args, **op_kwargs)
assert isinstance(output, PokeReturnValue)
Expand All @@ -259,7 +258,6 @@ def test_file_path_present_with_callback(self, sftp_hook_mock, op_args, op_kwarg
)
@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_file_pattern_present_with_callback(self, sftp_hook_mock, op_args, op_kwargs):
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sample_callable = Mock()
sample_callable.return_value = ["sample_return"]
sftp_hook_mock.return_value.get_files_by_pattern.return_value = [
Expand All @@ -285,3 +283,27 @@ def test_file_pattern_present_with_callback(self, sftp_hook_mock, op_args, op_kw
"files_found": ["/path/to/file/text_file.txt", "/path/to/file/another_text_file.txt"],
"decorator_return_value": ["sample_return"],
}

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_mod_time_called_when_newer_than_set(self, sftp_hook_mock):
sftp_hook_mock.return_value.isfile.return_value = True
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
tz = timezone("America/Toronto")
sftp_sensor = SFTPSensor(
task_id="unit_test",
path="/path/to/file/1970-01-01.txt",
newer_than=tz.convert(datetime(1960, 1, 2)),
)
context = {"ds": "1970-01-01"}
sftp_sensor.poke(context)
sftp_hook_mock.return_value.isfile.assert_called_once_with("/path/to/file/1970-01-01.txt")
sftp_hook_mock.return_value.get_mod_time.assert_called_once_with("/path/to/file/1970-01-01.txt")

@patch("airflow.providers.sftp.sensors.sftp.SFTPHook")
def test_mod_time_not_called_when_newer_than_not_set(self, sftp_hook_mock):
sftp_hook_mock.return_value.isfile.return_value = True
sftp_hook_mock.return_value.get_mod_time.return_value = "19700101000000"
sftp_sensor = SFTPSensor(task_id="unit_test", path="/path/to/file/1970-01-01.txt")
context = {"ds": "1970-01-01"}
sftp_sensor.poke(context)
sftp_hook_mock.return_value.get_mod_time.assert_not_called()