[airflow] branch main updated: Fixing MyPy issues in testa/jobs (#19998)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 8ca3eda  Fixing MyPy issues in testa/jobs (#19998)
8ca3eda is described below

commit 8ca3eda02c58ab7ff224e463bd8b7ae389ce4370
Author: Khalid Mammadov <xm...@hotmail.com>
AuthorDate: Sun Dec 5 23:24:21 2021 +0000

    Fixing MyPy issues in testa/jobs (#19998)
---
 tests/jobs/test_local_task_job.py | 7 ++++---
 tests/jobs/test_scheduler_job.py  | 8 +++++++-
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/tests/jobs/test_local_task_job.py b/tests/jobs/test_local_task_job.py
index 8a4df12..150dcb4 100644
--- a/tests/jobs/test_local_task_job.py
+++ b/tests/jobs/test_local_task_job.py
@@ -22,6 +22,7 @@ import signal
 import time
 import uuid
 from multiprocessing import Lock, Value
+from typing import List, Union
 from unittest import mock
 from unittest.mock import patch
 
@@ -864,10 +865,10 @@ def clean_db_helper():
 
 
 @pytest.mark.usefixtures("clean_db_helper")
-@pytest.mark.parametrize("return_codes", [[0], 9 * [None] + [0]])
 @mock.patch("airflow.jobs.local_task_job.get_task_runner")
-def test_number_of_queries_single_loop(mock_get_task_runner, return_codes, dag_maker):
-    mock_get_task_runner.return_value.return_code.side_effects = return_codes
+def test_number_of_queries_single_loop(mock_get_task_runner, dag_maker):
+    codes: List[Union[int, None]] = 9 * [None] + [0]
+    mock_get_task_runner.return_value.return_code.side_effects = [[0], codes]
 
     unique_prefix = str(uuid.uuid4())
     with dag_maker(dag_id=f'{unique_prefix}_test_number_of_queries'):
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 3653690..5f4d854 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -22,6 +22,7 @@ import os
 import shutil
 from datetime import timedelta
 from tempfile import mkdtemp
+from typing import Generator, Optional
 from unittest import mock
 from unittest.mock import MagicMock, patch
 
@@ -1897,6 +1898,7 @@ class TestSchedulerJob:
 
         def _create_dagruns(dag: DAG):
             next_info = dag.next_dagrun_info(None)
+            assert next_info is not None
             for _ in range(5):
                 yield dag.create_dagrun(
                     run_type=DagRunType.SCHEDULED,
@@ -1905,6 +1907,8 @@ class TestSchedulerJob:
                     state=State.RUNNING,
                 )
                 next_info = dag.next_dagrun_info(next_info.data_interval)
+                if next_info is None:
+                    break
 
         # Create 5 dagruns for each DAG.
         # To increase the chances the TIs from the "full" pool will get retrieved first, we schedule all
@@ -3435,6 +3439,8 @@ class TestSchedulerJobQueriesCount:
     made that affects the performance of the SchedulerJob.
     """
 
+    scheduler_job: Optional[SchedulerJob]
+
     @staticmethod
     def clean_db():
         clear_db_runs()
@@ -3446,7 +3452,7 @@ class TestSchedulerJobQueriesCount:
         clear_db_serialized_dags()
 
     @pytest.fixture(autouse=True)
-    def per_test(self) -> None:
+    def per_test(self) -> Generator:
         self.clean_db()
 
         yield