Migrated operators tests to pytest (#29377)

* Migrated operators tests to pytest

* migrated remaining operators tests to pytest

* resolved static checks

* refactored

* refactored

Author: Abhishek-kumar-samsung

tests/operators/test_bash.py

@@ -25,7 +25,6 @@
 from unittest import mock
 
 import pytest
-from parameterized import parameterized
 
 from airflow.exceptions import AirflowException, AirflowSkipException, AirflowTaskTimeout
 from airflow.models.dag import DAG
@@ -40,11 +39,12 @@
 
 
 class TestBashOperator:
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "append_env,user_defined_env,expected_airflow_home",
         [
             (False, None, "MY_PATH_TO_AIRFLOW_HOME"),
             (True, {"AIRFLOW_HOME": "OVERRIDDEN_AIRFLOW_HOME"}, "OVERRIDDEN_AIRFLOW_HOME"),
-        ]
+        ],
     )
     def test_echo_env_variables(self, append_env, user_defined_env, expected_airflow_home):
         """
@@ -98,13 +98,14 @@ def test_echo_env_variables(self, append_env, user_defined_env, expected_airflow
                 output = "".join(file.readlines())
                 assert expected == output
 
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "val,expected",
         [
             ("test-val", "test-val"),
             ("test-val\ntest-val\n", ""),
             ("test-val\ntest-val", "test-val"),
             ("", ""),
-        ]
+        ],
     )
     def test_return_value(self, val, expected):
         op = BashOperator(task_id="abc", bash_command=f'set -e; echo "{val}";')
@@ -168,13 +169,14 @@ def test_valid_cwd(self):
             with open(f"{test_cwd_folder}/outputs.txt") as tmp_file:
                 assert tmp_file.read().splitlines()[0] == "xxxx"
 
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "extra_kwargs,actual_exit_code,expected_exc",
         [
             (None, 99, AirflowSkipException),
             ({"skip_exit_code": 100}, 100, AirflowSkipException),
             ({"skip_exit_code": 100}, 101, AirflowException),
             ({"skip_exit_code": None}, 99, AirflowException),
-        ]
+        ],
     )
     def test_skip(self, extra_kwargs, actual_exit_code, expected_exc):
         kwargs = dict(task_id="abc", bash_command=f'set -e; echo "hello world"; exit {actual_exit_code};')

tests/operators/test_datetime.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import datetime
-import unittest
 
 import pytest
 import time_machine
@@ -35,24 +34,22 @@
 INTERVAL = datetime.timedelta(hours=12)
 
 
-class TestBranchDateTimeOperator(unittest.TestCase):
+class TestBranchDateTimeOperator:
     @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
+    def setup_class(cls):
 
         with create_session() as session:
             session.query(DagRun).delete()
             session.query(TI).delete()
 
-        cls.targets = [
-            (datetime.datetime(2020, 7, 7, 10, 0, 0), datetime.datetime(2020, 7, 7, 11, 0, 0)),
-            (datetime.time(10, 0, 0), datetime.time(11, 0, 0)),
-            (datetime.datetime(2020, 7, 7, 10, 0, 0), datetime.time(11, 0, 0)),
-            (datetime.time(10, 0, 0), datetime.datetime(2020, 7, 7, 11, 0, 0)),
-            (datetime.time(11, 0, 0), datetime.time(10, 0, 0)),
-        ]
+    targets = [
+        (datetime.datetime(2020, 7, 7, 10, 0, 0), datetime.datetime(2020, 7, 7, 11, 0, 0)),
+        (datetime.time(10, 0, 0), datetime.time(11, 0, 0)),
+        (datetime.datetime(2020, 7, 7, 10, 0, 0), datetime.time(11, 0, 0)),
+        (datetime.time(10, 0, 0), datetime.datetime(2020, 7, 7, 11, 0, 0)),
+    ]
 
-    def setUp(self):
+    def setup_method(self):
         self.dag = DAG(
             "branch_datetime_operator_test",
             default_args={"owner": "airflow", "start_date": DEFAULT_DATE},
@@ -79,8 +76,7 @@ def setUp(self):
             run_id="manual__", start_date=DEFAULT_DATE, execution_date=DEFAULT_DATE, state=State.RUNNING
         )
 
-    def tearDown(self):
-        super().tearDown()
+    def teardown_method(self):
 
         with create_session() as session:
             session.query(DagRun).delete()
@@ -95,15 +91,13 @@ def _assert_task_ids_match_states(self, task_ids_to_states):
             except KeyError:
                 raise ValueError(f"Invalid task id {ti.task_id} found!")
             else:
-                self.assertEqual(
-                    ti.state,
-                    expected_state,
-                    f"Task {ti.task_id} has state {ti.state} instead of expected {expected_state}",
-                )
+                assert (ti.state) == (
+                    expected_state
+                ), f"Task {ti.task_id} has state {ti.state} instead of expected {expected_state}"
 
     def test_no_target_time(self):
         """Check if BranchDateTimeOperator raises exception on missing target"""
-        with self.assertRaises(AirflowException):
+        with pytest.raises(AirflowException):
             BranchDateTimeOperator(
                 task_id="datetime_branch",
                 follow_task_ids_if_true="branch_1",
@@ -113,121 +107,125 @@ def test_no_target_time(self):
                 dag=self.dag,
             )
 
+    @pytest.mark.parametrize(
+        "target_lower,target_upper",
+        targets,
+    )
     @time_machine.travel("2020-07-07 10:54:05")
-    def test_branch_datetime_operator_falls_within_range(self):
+    def test_branch_datetime_operator_falls_within_range(self, target_lower, target_upper):
         """Check BranchDateTimeOperator branch operation"""
-        for target_lower, target_upper in self.targets:
-            with self.subTest(target_lower=target_lower, target_upper=target_upper):
-                self.branch_op.target_lower = target_lower
-                self.branch_op.target_upper = target_upper
-                self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
-
-                self._assert_task_ids_match_states(
-                    {
-                        "datetime_branch": State.SUCCESS,
-                        "branch_1": State.NONE,
-                        "branch_2": State.SKIPPED,
-                    }
-                )
+        self.branch_op.target_lower = target_lower
+        self.branch_op.target_upper = target_upper
+        self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+        self._assert_task_ids_match_states(
+            {
+                "datetime_branch": State.SUCCESS,
+                "branch_1": State.NONE,
+                "branch_2": State.SKIPPED,
+            }
+        )
 
-    def test_branch_datetime_operator_falls_outside_range(self):
+    @pytest.mark.parametrize(
+        "target_lower,target_upper",
+        targets,
+    )
+    def test_branch_datetime_operator_falls_outside_range(self, target_lower, target_upper):
         """Check BranchDateTimeOperator branch operation"""
         dates = [
             datetime.datetime(2020, 7, 7, 12, 0, 0, tzinfo=datetime.timezone.utc),
             datetime.datetime(2020, 6, 7, 12, 0, 0, tzinfo=datetime.timezone.utc),
         ]
 
-        for target_lower, target_upper in self.targets:
-            with self.subTest(target_lower=target_lower, target_upper=target_upper):
-                self.branch_op.target_lower = target_lower
-                self.branch_op.target_upper = target_upper
-
-                for date in dates:
-                    with time_machine.travel(date):
-                        self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
-
-                        self._assert_task_ids_match_states(
-                            {
-                                "datetime_branch": State.SUCCESS,
-                                "branch_1": State.SKIPPED,
-                                "branch_2": State.NONE,
-                            }
-                        )
-
-    @time_machine.travel("2020-07-07 10:54:05")
-    def test_branch_datetime_operator_upper_comparison_within_range(self):
-        """Check BranchDateTimeOperator branch operation"""
-        for _, target_upper in self.targets:
-            with self.subTest(target_upper=target_upper):
-                self.branch_op.target_upper = target_upper
-                self.branch_op.target_lower = None
+        self.branch_op.target_lower = target_lower
+        self.branch_op.target_upper = target_upper
 
+        for date in dates:
+            with time_machine.travel(date):
                 self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
 
                 self._assert_task_ids_match_states(
                     {
                         "datetime_branch": State.SUCCESS,
-                        "branch_1": State.NONE,
-                        "branch_2": State.SKIPPED,
+                        "branch_1": State.SKIPPED,
+                        "branch_2": State.NONE,
                     }
                 )
 
+    @pytest.mark.parametrize("target_upper", [target_upper for (_, target_upper) in targets])
     @time_machine.travel("2020-07-07 10:54:05")
-    def test_branch_datetime_operator_lower_comparison_within_range(self):
+    def test_branch_datetime_operator_upper_comparison_within_range(self, target_upper):
         """Check BranchDateTimeOperator branch operation"""
-        for target_lower, _ in self.targets:
-            with self.subTest(target_lower=target_lower):
-                self.branch_op.target_lower = target_lower
-                self.branch_op.target_upper = None
+        self.branch_op.target_upper = target_upper
+        self.branch_op.target_lower = None
 
-                self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+        self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
 
-                self._assert_task_ids_match_states(
-                    {
-                        "datetime_branch": State.SUCCESS,
-                        "branch_1": State.NONE,
-                        "branch_2": State.SKIPPED,
-                    }
-                )
+        self._assert_task_ids_match_states(
+            {
+                "datetime_branch": State.SUCCESS,
+                "branch_1": State.NONE,
+                "branch_2": State.SKIPPED,
+            }
+        )
+
+    @pytest.mark.parametrize("target_lower", [target_lower for (target_lower, _) in targets])
+    @time_machine.travel("2020-07-07 10:54:05")
+    def test_branch_datetime_operator_lower_comparison_within_range(self, target_lower):
+        """Check BranchDateTimeOperator branch operation"""
+        self.branch_op.target_lower = target_lower
+        self.branch_op.target_upper = None
 
+        self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+
+        self._assert_task_ids_match_states(
+            {
+                "datetime_branch": State.SUCCESS,
+                "branch_1": State.NONE,
+                "branch_2": State.SKIPPED,
+            }
+        )
+
+    @pytest.mark.parametrize("target_upper", [target_upper for (_, target_upper) in targets])
     @time_machine.travel("2020-07-07 12:00:00")
-    def test_branch_datetime_operator_upper_comparison_outside_range(self):
+    def test_branch_datetime_operator_upper_comparison_outside_range(self, target_upper):
         """Check BranchDateTimeOperator branch operation"""
-        for _, target_upper in self.targets:
-            with self.subTest(target_upper=target_upper):
-                self.branch_op.target_upper = target_upper
-                self.branch_op.target_lower = None
+        self.branch_op.target_upper = target_upper
+        self.branch_op.target_lower = None
 
-                self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+        self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
 
-                self._assert_task_ids_match_states(
-                    {
-                        "datetime_branch": State.SUCCESS,
-                        "branch_1": State.SKIPPED,
-                        "branch_2": State.NONE,
-                    }
-                )
+        self._assert_task_ids_match_states(
+            {
+                "datetime_branch": State.SUCCESS,
+                "branch_1": State.SKIPPED,
+                "branch_2": State.NONE,
+            }
+        )
 
+    @pytest.mark.parametrize("target_lower", [target_lower for (target_lower, _) in targets])
     @time_machine.travel("2020-07-07 09:00:00")
-    def test_branch_datetime_operator_lower_comparison_outside_range(self):
+    def test_branch_datetime_operator_lower_comparison_outside_range(self, target_lower):
         """Check BranchDateTimeOperator branch operation"""
-        for target_lower, _ in self.targets:
-            with self.subTest(target_lower=target_lower):
-                self.branch_op.target_lower = target_lower
-                self.branch_op.target_upper = None
+        self.branch_op.target_lower = target_lower
+        self.branch_op.target_upper = None
 
-                self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+        self.branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
 
-                self._assert_task_ids_match_states(
-                    {
-                        "datetime_branch": State.SUCCESS,
-                        "branch_1": State.SKIPPED,
-                        "branch_2": State.NONE,
-                    }
-                )
+        self._assert_task_ids_match_states(
+            {
+                "datetime_branch": State.SUCCESS,
+                "branch_1": State.SKIPPED,
+                "branch_2": State.NONE,
+            }
+        )
 
+    @pytest.mark.parametrize(
+        "target_lower,target_upper",
+        targets,
+    )
     @time_machine.travel("2020-12-01 09:00:00")
-    def test_branch_datetime_operator_use_task_logical_date(self):
+    def test_branch_datetime_operator_use_task_logical_date(self, target_lower, target_upper):
         """Check if BranchDateTimeOperator uses task execution date"""
         in_between_date = timezone.datetime(2020, 7, 7, 10, 30, 0)
         self.branch_op.use_task_logical_date = True
@@ -238,19 +236,17 @@ def test_branch_datetime_operator_use_task_logical_date(self):
             state=State.RUNNING,
         )
 
-        for target_lower, target_upper in self.targets:
-            with self.subTest(target_lower=target_lower, target_upper=target_upper):
-                self.branch_op.target_lower = target_lower
-                self.branch_op.target_upper = target_upper
-                self.branch_op.run(start_date=in_between_date, end_date=in_between_date)
+        self.branch_op.target_lower = target_lower
+        self.branch_op.target_upper = target_upper
+        self.branch_op.run(start_date=in_between_date, end_date=in_between_date)
 
-                self._assert_task_ids_match_states(
-                    {
-                        "datetime_branch": State.SUCCESS,
-                        "branch_1": State.NONE,
-                        "branch_2": State.SKIPPED,
-                    }
-                )
+        self._assert_task_ids_match_states(
+            {
+                "datetime_branch": State.SUCCESS,
+                "branch_1": State.NONE,
+                "branch_2": State.SKIPPED,
+            }
+        )
 
     def test_deprecation_warning(self):
         warning_message = (

tests/operators/test_email.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import datetime
-import unittest
 from unittest import mock
 
 from airflow.models.dag import DAG
@@ -34,15 +33,13 @@
 send_email_test = mock.Mock()
 
 
-class TestEmailOperator(unittest.TestCase):
-    def setUp(self):
-        super().setUp()
+class TestEmailOperator:
+    def setup_class(self):
         self.dag = DAG(
             "test_dag",
             default_args={"owner": "airflow", "start_date": DEFAULT_DATE},
             schedule=INTERVAL,
         )
-        self.addCleanup(self.dag.clear)
 
     def _run_as_operator(self, **kwargs):
         task = EmailOperator(
@@ -56,6 +53,7 @@ def _run_as_operator(self, **kwargs):
             **kwargs,
         )
         task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+        self.dag.clear()
 
     def test_execute(self):
         with conf_vars({("email", "email_backend"): "tests.operators.test_email.send_email_test"}):