[airflow] branch master updated: Don't set child tasks to schedulable in test runs (#12595)

[di...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new 6caf260  Don't set child tasks to schedulable in test runs (#12595)
6caf260 is described below

commit 6caf2607e04f581abdcb38fdbc426e03d5307429
Author: Daniel Imberman <da...@gmail.com>
AuthorDate: Tue Nov 24 13:43:22 2020 -0800

    Don't set child tasks to schedulable in test runs (#12595)
    
    Fixes a bug where Airflow will attempt to set child tasks to schedulable
    for test tasks when users run `airflow task test.` This causes an error
    as Airflow runs a DB seek for a task that has not been recorded.
---
 airflow/models/taskinstance.py          | 3 ++-
 tests/cli/commands/test_task_command.py | 5 ++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index c6a1378..2059542 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1149,7 +1149,8 @@ class TaskInstance(Base, LoggingMixin):  # pylint: disable=R0902,R0904
 
         session.commit()
 
-        self._run_mini_scheduler_on_child_tasks(session)
+        if not test_mode:
+            self._run_mini_scheduler_on_child_tasks(session)
 
     @provide_session
     @Sentry.enrich_errors
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
index 46e3d61..61020ec 100644
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -69,7 +69,8 @@ class TestCliTasks(unittest.TestCase):
         args = self.parser.parse_args(['tasks', 'list', 'example_bash_operator', '--tree'])
         task_command.task_list(args)
 
-    def test_test(self):
+    @mock.patch("airflow.models.taskinstance.TaskInstance._run_mini_scheduler_on_child_tasks")
+    def test_test(self, mock_run_mini_scheduler):
         """Test the `airflow test` command"""
         args = self.parser.parse_args(
             ["tasks", "test", "example_python_operator", 'print_the_context', '2018-01-01']
@@ -77,6 +78,8 @@ class TestCliTasks(unittest.TestCase):
 
         with redirect_stdout(io.StringIO()) as stdout:
             task_command.task_test(args)
+
+        mock_run_mini_scheduler.assert_not_called()
         # Check that prints, and log messages, are shown
         self.assertIn("'example_python_operator__print_the_context__20180101'", stdout.getvalue())