Fixing mypy issues inside tests model (apache#20026)

tests/models/test_baseoperator.py

@@ -17,6 +17,7 @@
 # under the License.
 import uuid
 from datetime import date, datetime
+from typing import Any
 from unittest import mock
 
 import jinja2
@@ -577,9 +578,11 @@ def test_weight_rule_override(self):
 
 def test_init_subclass_args():
     class InitSubclassOp(BaseOperator):
+        _class_arg: Any
+
         def __init_subclass__(cls, class_arg=None, **kwargs) -> None:
             cls._class_arg = class_arg
-            super().__init_subclass__(**kwargs)
+            super().__init_subclass__()
 
         def execute(self, context):
             self.context_arg = context

tests/models/test_dag.py

@@ -1514,8 +1514,8 @@ def test_clear_set_dagrun_state_for_parent_dag(self, dag_run_state):
 
     @parameterized.expand(
         [(state, State.NONE) for state in State.task_states if state != State.RUNNING]
-        + [(State.RUNNING, State.RESTARTING)]
-    )  # type: ignore
+        + [(State.RUNNING, State.RESTARTING)]  # type: ignore
+    )
     def test_clear_dag(self, ti_state_begin, ti_state_end: Optional[str]):
         dag_id = 'test_clear_dag'
         self._clean_up(dag_id)

tests/models/test_taskinstance.py

@@ -63,7 +63,7 @@
 from airflow.utils import timezone
 from airflow.utils.db import merge_conn
 from airflow.utils.session import create_session, provide_session
-from airflow.utils.state import State
+from airflow.utils.state import State, TaskInstanceState
 from airflow.utils.types import DagRunType
 from airflow.version import version
 from tests.models import DEFAULT_DATE, TEST_DAGS_FOLDER
@@ -972,7 +972,8 @@ def test_check_task_dependencies(
             for i in range(5):
                 task = DummyOperator(task_id=f'runme_{i}', dag=dag)
                 task.set_downstream(downstream)
-        run_date = task.start_date + datetime.timedelta(days=5)
+            assert task.start_date is not None
+            run_date = task.start_date + datetime.timedelta(days=5)
 
         ti = dag_maker.create_dagrun(execution_date=run_date).get_task_instance(downstream.task_id)
         ti.task = downstream
@@ -1374,7 +1375,10 @@ def test_success_callback_no_race_condition(self, create_task_instance):
 
     @staticmethod
     def _test_previous_dates_setup(
-        schedule_interval: Union[str, datetime.timedelta, None], catchup: bool, scenario: List[str], dag_maker
+        schedule_interval: Union[str, datetime.timedelta, None],
+        catchup: bool,
+        scenario: List[TaskInstanceState],
+        dag_maker,
     ) -> list:
         dag_id = 'test_previous_dates'
         with dag_maker(dag_id=dag_id, schedule_interval=schedule_interval, catchup=catchup):