From 7bb49abc6fcd973de50b279cd3b9d770776566a2 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Sun, 24 Oct 2021 12:17:21 -0400
Subject: [PATCH 01/11] Refactored tests. Ensured State.FAILED considered in
 one case.

---
 tests/models/test_taskinstance.py | 41 ++++++++++---------------------
 1 file changed, 13 insertions(+), 28 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index db6e9e384dfc6..0d701a0c17ce6 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -492,41 +492,26 @@ def test_task_wipes_next_fields(self, session, state, dag_maker):
         when they go into a state of FAILED, SKIPPED, SUCCESS, UP_FOR_RESCHEDULE, or UP_FOR_RETRY.
         """
 
-        def failure():
-            raise AirflowException
-
-        def skip():
-            raise AirflowSkipException
-
-        def success():
-            return None
-
-        def reschedule():
-            reschedule_date = timezone.utcnow()
-            raise AirflowRescheduleException(reschedule_date)
+        def run(state):
+            if state in [State.FAILED, State.UP_FOR_RETRY]:
+                raise AirflowException
+            if state == State.SKIPPED:
+                raise AirflowSkipException
+            if state == State.UP_FOR_RESCHEDULE:
+                raise AirflowRescheduleException(timezone.utcnow())
+            return None  # SUCCESS
 
         _retries = 0
-        _retry_delay = datetime.timedelta(seconds=0)
-
-        if state == State.FAILED:
-            _python_callable = failure
-        elif state == State.SKIPPED:
-            _python_callable = skip
-        elif state == State.SUCCESS:
-            _python_callable = success
-        elif state == State.UP_FOR_RESCHEDULE:
-            _python_callable = reschedule
-        elif state in [State.FAILED, State.UP_FOR_RETRY]:
-            _python_callable = failure
+        if state == State.UP_FOR_RETRY:
             _retries = 1
-            _retry_delay = datetime.timedelta(seconds=2)
 
         with dag_maker("test_deferred_method_clear"):
             task = PythonOperator(
                 task_id="test_deferred_method_clear_task",
-                python_callable=_python_callable,
+                python_callable=run,
+                op_args=[state],
                 retries=_retries,
-                retry_delay=_retry_delay,
+                retry_delay=datetime.timedelta(seconds=2),
             )
 
         dr = dag_maker.create_dagrun()
@@ -540,7 +525,7 @@ def reschedule():
         if state in [State.FAILED, State.UP_FOR_RETRY]:
             with pytest.raises(AirflowException):
                 ti.run()
-        elif state in [State.SKIPPED, State.SUCCESS, State.UP_FOR_RESCHEDULE]:
+        else:
             ti.run()
         ti.refresh_from_db()
 

From 0f84246f3552f1094e601c31cfd29c6f9ee1f2f2 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Sun, 24 Oct 2021 12:19:55 -0400
Subject: [PATCH 02/11] Updated Comment.

---
 tests/models/test_taskinstance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 0d701a0c17ce6..9a73f22b8ec62 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -499,7 +499,7 @@ def run(state):
                 raise AirflowSkipException
             if state == State.UP_FOR_RESCHEDULE:
                 raise AirflowRescheduleException(timezone.utcnow())
-            return None  # SUCCESS
+            return None  # State.SUCCESS
 
         _retries = 0
         if state == State.UP_FOR_RETRY:

From ea1049d0d0c67a0edd4ba2dd7e32caa8a4410685 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Mon, 25 Oct 2021 10:41:43 -0400
Subject: [PATCH 03/11] Refactored tests with further parametrization.

---
 tests/models/test_taskinstance.py | 47 ++++++++++++++++---------------
 1 file changed, 25 insertions(+), 22 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 9a73f22b8ec62..86afc89461ef7 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -141,6 +141,18 @@ def test_load_error_file_loads_correctly(self):
             set_error_file(error_fd.name, error=error_message)
             assert load_error_file(error_fd) == error_message
 
+    def _failure():
+        raise AirflowException
+
+    def _reschedule():
+        raise AirflowRescheduleException(timezone.utcnow())
+
+    def _skip():
+        raise AirflowSkipException
+
+    def _success():
+        return None
+
     def test_set_task_dates(self, dag_maker):
         """
         Test that tasks properly take start/end dates from DAGs
@@ -483,34 +495,25 @@ def task_function(ti):
         ti.state == state
 
     @pytest.mark.parametrize(
-        "state",
-        [State.FAILED, State.SKIPPED, State.SUCCESS, State.UP_FOR_RESCHEDULE, State.UP_FOR_RETRY],
+        "state, func, retries",
+        [
+            (State.FAILED, _failure, 0),
+            (State.SKIPPED, _skip, 0),
+            (State.SUCCESS, _success, 0),
+            (State.UP_FOR_RESCHEDULE, _reschedule, 0),
+            (State.UP_FOR_RETRY, _failure, 1),
+        ],
     )
-    def test_task_wipes_next_fields(self, session, state, dag_maker):
+    def test_task_wipes_next_fields(self, session, dag_maker, state, func, retries):
         """
-        Test that ensures that tasks wipe their next_method and next_kwargs
-        when they go into a state of FAILED, SKIPPED, SUCCESS, UP_FOR_RESCHEDULE, or UP_FOR_RETRY.
+        Test that ensures that tasks wipe their next_method and next_kwargs for the configured states.
         """
 
-        def run(state):
-            if state in [State.FAILED, State.UP_FOR_RETRY]:
-                raise AirflowException
-            if state == State.SKIPPED:
-                raise AirflowSkipException
-            if state == State.UP_FOR_RESCHEDULE:
-                raise AirflowRescheduleException(timezone.utcnow())
-            return None  # State.SUCCESS
-
-        _retries = 0
-        if state == State.UP_FOR_RETRY:
-            _retries = 1
-
         with dag_maker("test_deferred_method_clear"):
             task = PythonOperator(
                 task_id="test_deferred_method_clear_task",
-                python_callable=run,
-                op_args=[state],
-                retries=_retries,
+                python_callable=func,
+                retries=retries,
                 retry_delay=datetime.timedelta(seconds=2),
             )
 
@@ -522,7 +525,7 @@ def run(state):
         session.commit()
 
         ti.task = task
-        if state in [State.FAILED, State.UP_FOR_RETRY]:
+        if func.__name__ == "_failure":
             with pytest.raises(AirflowException):
                 ti.run()
         else:

From adef6111b6d9544b1159dec41bebee6ce1b6002a Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Mon, 25 Oct 2021 12:08:32 -0400
Subject: [PATCH 04/11] Updated failure function check.

---
 tests/models/test_taskinstance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 86afc89461ef7..5a0e182db4696 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -525,7 +525,7 @@ def test_task_wipes_next_fields(self, session, dag_maker, state, func, retries):
         session.commit()
 
         ti.task = task
-        if func.__name__ == "_failure":
+        if func is type(self)._failure:
             with pytest.raises(AirflowException):
                 ti.run()
         else:

From 7278184c7f98103bc37fd26e96a9e1d65c82f13a Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Thu, 23 Dec 2021 16:04:26 -0500
Subject: [PATCH 05/11] Updated tests for mypy error: Method must have at least
 one argument.

---
 tests/models/test_taskinstance.py | 35 +++++++++++++------------------
 1 file changed, 14 insertions(+), 21 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 5a0e182db4696..9d3e27d20b41e 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -141,18 +141,6 @@ def test_load_error_file_loads_correctly(self):
             set_error_file(error_fd.name, error=error_message)
             assert load_error_file(error_fd) == error_message
 
-    def _failure():
-        raise AirflowException
-
-    def _reschedule():
-        raise AirflowRescheduleException(timezone.utcnow())
-
-    def _skip():
-        raise AirflowSkipException
-
-    def _success():
-        return None
-
     def test_set_task_dates(self, dag_maker):
         """
         Test that tasks properly take start/end dates from DAGs
@@ -495,24 +483,29 @@ def task_function(ti):
         ti.state == state
 
     @pytest.mark.parametrize(
-        "state, func, retries",
+        "state, exception_type, retries",
         [
-            (State.FAILED, _failure, 0),
-            (State.SKIPPED, _skip, 0),
-            (State.SUCCESS, _success, 0),
-            (State.UP_FOR_RESCHEDULE, _reschedule, 0),
-            (State.UP_FOR_RETRY, _failure, 1),
+            (State.FAILED, AirflowException, 0),
+            (State.SKIPPED, AirflowSkipException, 0),
+            (State.SUCCESS, None, 0),
+            (State.UP_FOR_RESCHEDULE, AirflowRescheduleException(timezone.utcnow()), 0),
+            (State.UP_FOR_RETRY, AirflowException, 1),
         ],
     )
-    def test_task_wipes_next_fields(self, session, dag_maker, state, func, retries):
+    def test_task_wipes_next_fields(self, session, dag_maker, state, exception_type, retries):
         """
         Test that ensures that tasks wipe their next_method and next_kwargs for the configured states.
         """
 
+        def _raise_af_exception(exception_type):
+            if exception_type:
+                raise exception_type
+
         with dag_maker("test_deferred_method_clear"):
             task = PythonOperator(
                 task_id="test_deferred_method_clear_task",
-                python_callable=func,
+                python_callable=_raise_af_exception,
+                op_args=[exception_type],
                 retries=retries,
                 retry_delay=datetime.timedelta(seconds=2),
             )
@@ -525,7 +518,7 @@ def test_task_wipes_next_fields(self, session, dag_maker, state, func, retries):
         session.commit()
 
         ti.task = task
-        if func is type(self)._failure:
+        if exception_type == AirflowException:
             with pytest.raises(AirflowException):
                 ti.run()
         else:

From 3dd9f733a29e98d55fea722d0a4c90bfd4b540c9 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Sun, 24 Oct 2021 12:17:21 -0400
Subject: [PATCH 06/11] Refactored tests. Ensured State.FAILED considered in
 one case.

---
 tests/models/test_taskinstance.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 9d3e27d20b41e..dad05224749bd 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -497,16 +497,37 @@ def test_task_wipes_next_fields(self, session, dag_maker, state, exception_type,
         Test that ensures that tasks wipe their next_method and next_kwargs for the configured states.
         """
 
+<<<<<<< HEAD
         def _raise_af_exception(exception_type):
             if exception_type:
                 raise exception_type
+=======
+        def run(state):
+            if state in [State.FAILED, State.UP_FOR_RETRY]:
+                raise AirflowException
+            if state == State.SKIPPED:
+                raise AirflowSkipException
+            if state == State.UP_FOR_RESCHEDULE:
+                raise AirflowRescheduleException(timezone.utcnow())
+            return None  # SUCCESS
+
+        _retries = 0
+        if state == State.UP_FOR_RETRY:
+            _retries = 1
+>>>>>>> 473345a70 (Refactored tests. Ensured State.FAILED considered in one case.)
 
         with dag_maker("test_deferred_method_clear"):
             task = PythonOperator(
                 task_id="test_deferred_method_clear_task",
+<<<<<<< HEAD
                 python_callable=_raise_af_exception,
                 op_args=[exception_type],
                 retries=retries,
+=======
+                python_callable=run,
+                op_args=[state],
+                retries=_retries,
+>>>>>>> 473345a70 (Refactored tests. Ensured State.FAILED considered in one case.)
                 retry_delay=datetime.timedelta(seconds=2),
             )
 

From 82317d9207163c71290331dbc885ddd70a53cbc3 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Sun, 24 Oct 2021 12:19:55 -0400
Subject: [PATCH 07/11] Updated Comment.

---
 tests/models/test_taskinstance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index dad05224749bd..dffb85669d3e0 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -509,7 +509,7 @@ def run(state):
                 raise AirflowSkipException
             if state == State.UP_FOR_RESCHEDULE:
                 raise AirflowRescheduleException(timezone.utcnow())
-            return None  # SUCCESS
+            return None  # State.SUCCESS
 
         _retries = 0
         if state == State.UP_FOR_RETRY:

From 99a32e495ef34976e31b466624378ebc65bf0c7c Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Thu, 23 Dec 2021 16:54:43 -0500
Subject: [PATCH 08/11] Updated tests for mypy error: Method must have at least
 one argument.

---
 tests/models/test_taskinstance.py | 37 +++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index dffb85669d3e0..b19c98497842c 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -141,6 +141,18 @@ def test_load_error_file_loads_correctly(self):
             set_error_file(error_fd.name, error=error_message)
             assert load_error_file(error_fd) == error_message
 
+    def _failure():
+        raise AirflowException
+
+    def _reschedule():
+        raise AirflowRescheduleException(timezone.utcnow())
+
+    def _skip():
+        raise AirflowSkipException
+
+    def _success():
+        return None
+
     def test_set_task_dates(self, dag_maker):
         """
         Test that tasks properly take start/end dates from DAGs
@@ -483,6 +495,7 @@ def task_function(ti):
         ti.state == state
 
     @pytest.mark.parametrize(
+<<<<<<< HEAD
         "state, exception_type, retries",
         [
             (State.FAILED, AirflowException, 0),
@@ -493,10 +506,23 @@ def task_function(ti):
         ],
     )
     def test_task_wipes_next_fields(self, session, dag_maker, state, exception_type, retries):
+=======
+        "state, func, retries",
+        [
+            (State.FAILED, _failure, 0),
+            (State.SKIPPED, _skip, 0),
+            (State.SUCCESS, _success, 0),
+            (State.UP_FOR_RESCHEDULE, _reschedule, 0),
+            (State.UP_FOR_RETRY, _failure, 1),
+        ],
+    )
+    def test_task_wipes_next_fields(self, session, dag_maker, state, func, retries):
+>>>>>>> 3ff2b4336 (Refactored tests with further parametrization.)
         """
         Test that ensures that tasks wipe their next_method and next_kwargs for the configured states.
         """
 
+<<<<<<< HEAD
 <<<<<<< HEAD
         def _raise_af_exception(exception_type):
             if exception_type:
@@ -528,6 +554,13 @@ def run(state):
                 op_args=[state],
                 retries=_retries,
 >>>>>>> 473345a70 (Refactored tests. Ensured State.FAILED considered in one case.)
+=======
+        with dag_maker("test_deferred_method_clear"):
+            task = PythonOperator(
+                task_id="test_deferred_method_clear_task",
+                python_callable=func,
+                retries=retries,
+>>>>>>> 3ff2b4336 (Refactored tests with further parametrization.)
                 retry_delay=datetime.timedelta(seconds=2),
             )
 
@@ -539,7 +572,11 @@ def run(state):
         session.commit()
 
         ti.task = task
+<<<<<<< HEAD
         if exception_type == AirflowException:
+=======
+        if func.__name__ == "_failure":
+>>>>>>> 3ff2b4336 (Refactored tests with further parametrization.)
             with pytest.raises(AirflowException):
                 ti.run()
         else:

From 130ccf82d23f5199ccdbbb89dd266a69cbf60ab9 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Thu, 23 Dec 2021 16:55:15 -0500
Subject: [PATCH 09/11] Updated tests for mypy error: Method must have at least
 one argument.

---
 tests/models/test_taskinstance.py | 58 -------------------------------
 1 file changed, 58 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index b19c98497842c..9d3e27d20b41e 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -141,18 +141,6 @@ def test_load_error_file_loads_correctly(self):
             set_error_file(error_fd.name, error=error_message)
             assert load_error_file(error_fd) == error_message
 
-    def _failure():
-        raise AirflowException
-
-    def _reschedule():
-        raise AirflowRescheduleException(timezone.utcnow())
-
-    def _skip():
-        raise AirflowSkipException
-
-    def _success():
-        return None
-
     def test_set_task_dates(self, dag_maker):
         """
         Test that tasks properly take start/end dates from DAGs
@@ -495,7 +483,6 @@ def task_function(ti):
         ti.state == state
 
     @pytest.mark.parametrize(
-<<<<<<< HEAD
         "state, exception_type, retries",
         [
             (State.FAILED, AirflowException, 0),
@@ -506,61 +493,20 @@ def task_function(ti):
         ],
     )
     def test_task_wipes_next_fields(self, session, dag_maker, state, exception_type, retries):
-=======
-        "state, func, retries",
-        [
-            (State.FAILED, _failure, 0),
-            (State.SKIPPED, _skip, 0),
-            (State.SUCCESS, _success, 0),
-            (State.UP_FOR_RESCHEDULE, _reschedule, 0),
-            (State.UP_FOR_RETRY, _failure, 1),
-        ],
-    )
-    def test_task_wipes_next_fields(self, session, dag_maker, state, func, retries):
->>>>>>> 3ff2b4336 (Refactored tests with further parametrization.)
         """
         Test that ensures that tasks wipe their next_method and next_kwargs for the configured states.
         """
 
-<<<<<<< HEAD
-<<<<<<< HEAD
         def _raise_af_exception(exception_type):
             if exception_type:
                 raise exception_type
-=======
-        def run(state):
-            if state in [State.FAILED, State.UP_FOR_RETRY]:
-                raise AirflowException
-            if state == State.SKIPPED:
-                raise AirflowSkipException
-            if state == State.UP_FOR_RESCHEDULE:
-                raise AirflowRescheduleException(timezone.utcnow())
-            return None  # State.SUCCESS
-
-        _retries = 0
-        if state == State.UP_FOR_RETRY:
-            _retries = 1
->>>>>>> 473345a70 (Refactored tests. Ensured State.FAILED considered in one case.)
 
         with dag_maker("test_deferred_method_clear"):
             task = PythonOperator(
                 task_id="test_deferred_method_clear_task",
-<<<<<<< HEAD
                 python_callable=_raise_af_exception,
                 op_args=[exception_type],
                 retries=retries,
-=======
-                python_callable=run,
-                op_args=[state],
-                retries=_retries,
->>>>>>> 473345a70 (Refactored tests. Ensured State.FAILED considered in one case.)
-=======
-        with dag_maker("test_deferred_method_clear"):
-            task = PythonOperator(
-                task_id="test_deferred_method_clear_task",
-                python_callable=func,
-                retries=retries,
->>>>>>> 3ff2b4336 (Refactored tests with further parametrization.)
                 retry_delay=datetime.timedelta(seconds=2),
             )
 
@@ -572,11 +518,7 @@ def run(state):
         session.commit()
 
         ti.task = task
-<<<<<<< HEAD
         if exception_type == AirflowException:
-=======
-        if func.__name__ == "_failure":
->>>>>>> 3ff2b4336 (Refactored tests with further parametrization.)
             with pytest.raises(AirflowException):
                 ti.run()
         else:

From 05d0adcc44061bf1136b5840150e0d48bbdb751f Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Thu, 23 Dec 2021 16:44:44 -0500
Subject: [PATCH 10/11] Added comment

---
 tests/models/test_taskinstance.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 9d3e27d20b41e..866a33a5a7aa9 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -494,7 +494,8 @@ def task_function(ti):
     )
     def test_task_wipes_next_fields(self, session, dag_maker, state, exception_type, retries):
         """
-        Test that ensures that tasks wipe their next_method and next_kwargs for the configured states.
+        Test that ensures that tasks wipe their next_method and next_kwargs for the configured states:
+        FAILED, SKIPPED, SUCCESS, UP_FOR_RESCHEDULE, UP_FOR_RETRY.
         """
 
         def _raise_af_exception(exception_type):

From 9b989196859038eac82f9a88aba958cefcacd451 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Thu, 6 Jan 2022 11:56:11 -0500
Subject: [PATCH 11/11] Renamed exception_type to exception, removed op_args
 and added exception directly to _raise_if_exception, and checked state vs.
 exception when running pytest.raises.

---
 tests/models/test_taskinstance.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 866a33a5a7aa9..fc741a226a396 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -483,7 +483,7 @@ def task_function(ti):
         ti.state == state
 
     @pytest.mark.parametrize(
-        "state, exception_type, retries",
+        "state, exception, retries",
         [
             (State.FAILED, AirflowException, 0),
             (State.SKIPPED, AirflowSkipException, 0),
@@ -492,21 +492,20 @@ def task_function(ti):
             (State.UP_FOR_RETRY, AirflowException, 1),
         ],
     )
-    def test_task_wipes_next_fields(self, session, dag_maker, state, exception_type, retries):
+    def test_task_wipes_next_fields(self, session, dag_maker, state, exception, retries):
         """
-        Test that ensures that tasks wipe their next_method and next_kwargs for the configured states:
-        FAILED, SKIPPED, SUCCESS, UP_FOR_RESCHEDULE, UP_FOR_RETRY.
+        Test that ensures that tasks wipe their next_method and next_kwargs
+        when the TI enters one of the configured states.
         """
 
-        def _raise_af_exception(exception_type):
-            if exception_type:
-                raise exception_type
+        def _raise_if_exception():
+            if exception:
+                raise exception
 
         with dag_maker("test_deferred_method_clear"):
             task = PythonOperator(
                 task_id="test_deferred_method_clear_task",
-                python_callable=_raise_af_exception,
-                op_args=[exception_type],
+                python_callable=_raise_if_exception,
                 retries=retries,
                 retry_delay=datetime.timedelta(seconds=2),
             )
@@ -519,8 +518,8 @@ def _raise_af_exception(exception_type):
         session.commit()
 
         ti.task = task
-        if exception_type == AirflowException:
-            with pytest.raises(AirflowException):
+        if state in [State.FAILED, State.UP_FOR_RETRY]:
+            with pytest.raises(exception):
                 ti.run()
         else:
             ti.run()
