RNHTTR commented on code in PR #41557:
URL: https://github.com/apache/airflow/pull/41557#discussion_r1723492253
##########
tests/sensors/test_time_delta.py:
##########
@@ -71,3 +71,17 @@ def test_timedelta_sensor(self, defer_mock, should_defer):
defer_mock.assert_called_once()
else:
defer_mock.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "should_defer",
+ [False, True],
+ )
+ @mock.patch("airflow.models.baseoperator.BaseOperator.defer")
+ def test_wait_sensor(self, defer_mock, should_defer):
+ delta = timedelta(seconds=30)
+ op = WaitSensor(task_id="wait_sensor_check", delta=delta,
dag=self.dag, deferrable=should_defer)
+ op.execute({})
+ if should_defer:
+ defer_mock.assert_called_once()
Review Comment:
Can you use `assert_called_once_with(...)`?
##########
tests/sensors/test_time_delta.py:
##########
@@ -71,3 +71,17 @@ def test_timedelta_sensor(self, defer_mock, should_defer):
defer_mock.assert_called_once()
else:
defer_mock.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "should_defer",
+ [False, True],
+ )
+ @mock.patch("airflow.models.baseoperator.BaseOperator.defer")
+ def test_wait_sensor(self, defer_mock, should_defer):
+ delta = timedelta(seconds=30)
Review Comment:
I wonder if you can patch `asyncio.sleep` so the test doesn't have to wait
at all
##########
airflow/sensors/time_delta.py:
##########
@@ -89,3 +92,38 @@ def execute(self, context: Context) -> bool | NoReturn:
def execute_complete(self, context: Context, event: Any = None) -> None:
"""Handle the event when the trigger fires and return immediately."""
return None
+
+
+class WaitSensor(BaseSensorOperator):
+ """
+ A sensor that waits a specified period of time before completing.
+
+ This differs from TimeDeltaSensor because the time to wait is measured
from the start of the task, not
+ the data_interval_end of the DAG run.
+
+ :param delta: time length to wait after the task starts before succeeding.
Review Comment:
What about `time_to_wait` instead of `delta`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
