Merge pull request #38 from airflow-laminar/tkp/tests

Add tests, fixes #22

Author: timkpaine

airflow_ha/__init__.py

@@ -2,4 +2,3 @@
 
 from .common import *
 from .operator import *
-from .utils import *

airflow_ha/operator.py

@@ -6,9 +6,9 @@
 from airflow.operators.python import BranchPythonOperator, PythonOperator
 from airflow.operators.trigger_dagrun import TriggerDagRunOperator
 from airflow.sensors.python import PythonSensor
+from airflow_common_operators import fail, pass_
 
 from .common import Action, CheckResult, Result
-from .utils import fail_, pass_
 
 __all__ = ("HighAvailabilityOperator",)
 _log = getLogger(__name__)
@@ -114,10 +114,10 @@ def __init__(
 
         # this is needed to ensure the dag fails, since the
         # retrigger_fail step will pass (to ensure dag retriggers!)
-        self._fail = PythonOperator(task_id=f"{self.task_id}-force-dag-fail", python_callable=fail_)
+        self._fail = PythonOperator(task_id=f"{self.task_id}-force-dag-fail", python_callable=fail)
 
         self._stop_pass = PythonOperator(task_id=f"{self.task_id}-stop-pass", python_callable=pass_)
-        self._stop_fail = PythonOperator(task_id=f"{self.task_id}-stop-fail", python_callable=fail_)
+        self._stop_fail = PythonOperator(task_id=f"{self.task_id}-stop-fail", python_callable=fail)
 
         # Update the retrigger counts in trigger kwargs
         retrigger_count_conf = f'''{{{{ (ti.dag_run.conf.get("{self.task_id}-retrigger", 0)|int) + 1 }}}}'''
@@ -175,6 +175,10 @@ def __init__(
         self._decide_task >> self._retrigger_pass
         self._decide_task >> self._retrigger_fail >> self._fail
 
+    @property
+    def decide_task(self) -> PythonOperator:
+        return self._decide_task
+
     @property
     def stop_fail(self) -> PythonOperator:
         return self._stop_fail
@@ -238,7 +242,7 @@ def _check_end_conditions(task_id, runtime, endtime, maxretrigger, start_date_or
             _log.info(
                 f"airflow-ha configuration -- endtime: {endtime}, endtime_as_datetime: {endtime_as_datetime}, datetime.now(tz=UTC): {datetime.now(tz=UTC)}"
             )
-            if endtime_as_datetime < datetime.now(tz=UTC):
+            if endtime_as_datetime <= datetime.now(tz=UTC):
                 # Endtime has passed, end
                 _log.info(f"endtime passed for {task_id}, stopping")
                 return Result.PASS, Action.STOP

airflow_ha/tests/conftest.py

@@ -0,0 +1,12 @@
+from airflow.models import DAG
+from pytest import fixture
+
+from airflow_ha import Action, HighAvailabilityOperator, Result
+
+
+@fixture(autouse=True)
+def operator():
+    callable = lambda **kwargs: (Result.PASS, Action.CONTINUE)  # noqa: E731
+    dag = DAG(dag_id="test_dag", default_args={}, schedule=None, params={})
+    operator = HighAvailabilityOperator(task_id="test_task", python_callable=callable, dag=dag)
+    return operator

airflow_ha/tests/test_operator.py

@@ -0,0 +1,277 @@
+from datetime import UTC, datetime, timedelta
+from unittest.mock import MagicMock, patch
+
+from airflow_ha import Action, HighAvailabilityOperator, Result
+from airflow_ha.operator import _callable_wrapper, _choose_branch
+
+
+class TestHighAvailabilityOperator:
+    def test_instantiation(self, operator: HighAvailabilityOperator):
+        assert operator.upstream_list == []
+        assert operator.downstream_list == [operator.decide_task]
+        assert operator.retrigger_fail in operator.decide_task.downstream_list
+        assert operator.retrigger_pass in operator.decide_task.downstream_list
+        assert operator.stop_pass in operator.decide_task.downstream_list
+        assert operator.stop_fail in operator.decide_task.downstream_list
+        assert operator.trigger_rule == "none_failed"
+
+    def test_check_end_conditions_default(self, operator: HighAvailabilityOperator):
+        dag_run_mock = MagicMock()
+        params_mock = MagicMock()
+        dag_run_mock.conf.get.side_effect = [
+            # airflow_ha_force_run
+            False,
+            # startdate
+            datetime(2025, 1, 1, tzinfo=UTC).isoformat(),
+            # retrigger
+            0,
+        ]
+        params_mock.get.side_effect = [
+            # force-run
+            False,
+            # force-runtime
+            None,
+            # force-endtime
+            None,
+            # maxretrigger
+            None,
+        ]
+        assert operator.check_end_conditions(dag_run=dag_run_mock, params=params_mock) is None
+
+    def test_check_end_conditions_force(self, operator: HighAvailabilityOperator):
+        dag_run_mock = MagicMock()
+        params_mock = MagicMock()
+        dag_run_mock.conf.get.side_effect = [
+            # airflow_ha_force_run
+            True,
+            # startdate
+            datetime(2025, 1, 1, tzinfo=UTC).isoformat(),
+            # retrigger
+            0,
+        ]
+        params_mock.get.side_effect = [
+            # force-run
+            False,
+            # force-runtime
+            None,
+            # force-endtime
+            None,
+            # maxretrigger
+            None,
+        ]
+        assert operator.check_end_conditions(dag_run=dag_run_mock, params=params_mock) is None
+
+    def test_check_end_conditions_runtime(self, operator: HighAvailabilityOperator):
+        now = datetime.now(tz=UTC)
+        yesterday_now = now - timedelta(days=1)
+        yesterday_now_almost = now - timedelta(hours=23, minutes=59, seconds=59)
+
+        dag_run_mock = MagicMock()
+        params_mock = MagicMock()
+        dag_run_mock.conf.get.side_effect = [
+            # airflow_ha_force_run
+            False,
+            # startdate
+            yesterday_now_almost.isoformat(),
+            # retrigger
+            0,
+            # airflow_ha_force_run
+            False,
+            # startdate
+            yesterday_now.isoformat(),
+            # retrigger
+            0,
+        ]
+        params_mock.get.side_effect = [
+            # force-run
+            False,
+            # force-runtime
+            timedelta(days=1),
+            # force-endtime
+            None,
+            # maxretrigger
+            None,
+            # force-run
+            False,
+            # force-runtime
+            timedelta(days=1),
+            # force-endtime
+            None,
+            # maxretrigger
+            None,
+        ]
+        assert operator.check_end_conditions(dag_run=dag_run_mock, params=params_mock) is None
+        assert operator.check_end_conditions(dag_run=dag_run_mock, params=params_mock) == (Result.PASS, Action.STOP)
+
+    def test_check_end_conditions_endtime(self, operator: HighAvailabilityOperator):
+        with patch("airflow_ha.operator.datetime") as mock_datetime:
+            now = datetime.now(tz=UTC)
+            one_second_ago = now - timedelta(seconds=1)
+            mock_datetime.combine = datetime.combine
+            mock_datetime.fromisoformat = datetime.fromisoformat
+            mock_datetime.now.return_value = one_second_ago
+
+            dag_run_mock = MagicMock()
+            params_mock = MagicMock()
+            dag_run_mock.conf.get.side_effect = [
+                # airflow_ha_force_run
+                False,
+                # startdate
+                one_second_ago.isoformat(),
+                # retrigger
+                0,
+                # airflow_ha_force_run
+                False,
+                # startdate
+                one_second_ago.isoformat(),
+                # retrigger
+                0,
+            ]
+            params_mock.get.side_effect = [
+                # force-run
+                False,
+                # force-runtime
+                None,
+                # force-endtime
+                now.time(),
+                # maxretrigger
+                None,
+                # force-run
+                False,
+                # force-runtime
+                None,
+                # force-endtime
+                now.time(),
+                # maxretrigger
+                None,
+            ]
+            assert operator.check_end_conditions(dag=operator.dag, dag_run=dag_run_mock, params=params_mock) is None
+            mock_datetime.now.return_value = now
+            assert operator.check_end_conditions(dag=operator.dag, dag_run=dag_run_mock, params=params_mock) == (Result.PASS, Action.STOP)
+
+    def test_check_end_conditions_maxretrigger(self, operator: HighAvailabilityOperator):
+        dag_run_mock = MagicMock()
+        params_mock = MagicMock()
+        dag_run_mock.conf.get.side_effect = [
+            # airflow_ha_force_run
+            False,
+            # startdate
+            None,
+            # retrigger
+            0,
+            # airflow_ha_force_run
+            False,
+            # startdate
+            None,
+            # retrigger
+            1,
+        ]
+        params_mock.get.side_effect = [
+            # force-run
+            False,
+            # force-runtime
+            None,
+            # force-endtime
+            None,
+            # maxretrigger
+            1,
+            # force-run
+            False,
+            # force-runtime
+            None,
+            # force-endtime
+            None,
+            # maxretrigger
+            1,
+        ]
+        assert operator.check_end_conditions(dag_run=dag_run_mock, params=params_mock) is None
+        assert operator.check_end_conditions(dag_run=dag_run_mock, params=params_mock) == (None, Action.STOP)
+
+    def test_choose_branch(self):
+        task_instance_mock = MagicMock()
+        branch_choices = {
+            (Result.PASS, Action.RETRIGGER): "a",
+            (Result.PASS, Action.STOP): "b",
+            (Result.FAIL, Action.RETRIGGER): "c",
+            (Result.FAIL, Action.STOP): "d",
+        }
+
+        # default
+        res = _choose_branch(
+            task_instance=task_instance_mock, branch_choices=branch_choices, task_id="test_task", check_end_conditions=lambda **kwargs: None
+        )
+        assert res == "a"
+
+        # pass, stop
+        task_instance_mock.xcom_pull.return_value = (Result.PASS, Action.STOP)
+        res = _choose_branch(
+            task_instance=task_instance_mock, branch_choices=branch_choices, task_id="test_task", check_end_conditions=lambda **kwargs: None
+        )
+        assert res == "b"
+        task_instance_mock.xcom_pull.return_value = (Result.FAIL, Action.RETRIGGER)
+        res = _choose_branch(
+            task_instance=task_instance_mock, branch_choices=branch_choices, task_id="test_task", check_end_conditions=lambda **kwargs: None
+        )
+        assert res == "c"
+        task_instance_mock.xcom_pull.return_value = (Result.FAIL, Action.STOP)
+        res = _choose_branch(
+            task_instance=task_instance_mock, branch_choices=branch_choices, task_id="test_task", check_end_conditions=lambda **kwargs: None
+        )
+        assert res == "d"
+
+        # end conditions met
+        res = _choose_branch(
+            task_instance=task_instance_mock,
+            branch_choices=branch_choices,
+            task_id="test_task",
+            check_end_conditions=lambda **kwargs: (Result.PASS, Action.STOP),
+        )
+        assert res == "b"
+
+        # retrigger exceeded
+        res = _choose_branch(
+            task_instance=task_instance_mock,
+            branch_choices=branch_choices,
+            task_id="test_task",
+            check_end_conditions=lambda **kwargs: (None, Action.STOP),
+        )
+        assert res == "b"
+
+    def test_callable_wrapper(self):
+        task_instance_mock = MagicMock()
+
+        # malformed
+        res = _callable_wrapper(
+            task_instance=task_instance_mock, python_callable=lambda **kwargs: (None, Action.CONTINUE), check_end_conditions=lambda **kwargs: None
+        )
+        assert res is True
+
+        # pass and continue
+        res = _callable_wrapper(
+            task_instance=task_instance_mock,
+            python_callable=lambda **kwargs: (Result.PASS, Action.CONTINUE),
+            check_end_conditions=lambda **kwargs: None,
+        )
+        assert res is False
+
+        # pass and retrigger (keep testing)
+        res = _callable_wrapper(
+            task_instance=task_instance_mock,
+            python_callable=lambda **kwargs: (Result.PASS, Action.RETRIGGER),
+            check_end_conditions=lambda **kwargs: None,
+        )
+        assert res is True
+
+        # pass and stop
+        res = _callable_wrapper(
+            task_instance=task_instance_mock, python_callable=lambda **kwargs: (Result.PASS, Action.STOP), check_end_conditions=lambda **kwargs: None
+        )
+        assert res is True
+
+        # end conditions met
+        res = _callable_wrapper(
+            task_instance=task_instance_mock,
+            python_callable=lambda **kwargs: (None, Action.CONTINUE),
+            check_end_conditions=lambda **kwargs: (None, Action.STOP),
+        )
+        assert res is True

airflow_ha/utils.py

@@ -30,6 +30,7 @@ classifiers = [
 
 dependencies = [
     "apache-airflow>=2.8,<3",
+    "airflow-common-operators>=0.2.0,<0.3",
 ]
 
 [project.optional-dependencies]
@@ -85,7 +86,7 @@ exclude_also = [
     "@(abc\\.)?abstractmethod",
 ]
 ignore_errors = true
-fail_under = 25
+fail_under = 90
 
 [tool.hatch.build]
 artifacts = []