def test_gauge_executor_metrics(self, mock_stats_gauge, mock_trigger_tasks, mock_sync):
executor = LocalExecutor()
executor.heartbeat()
calls = [mock.call('executor.open_slots', mock.ANY),
mock.call('executor.queued_tasks', mock.ANY),
mock.call('executor.running_tasks', mock.ANY)]
mock_stats_gauge.assert_has_calls(calls)
def test_scheduler_task(self):
TEST_DAG_FOLDER = os.environ['AIRFLOW__CORE__DAGS_FOLDER']
DEFAULT_DATE = timezone.datetime(2020, 1, 1)
dag_id = 'test_event_based_dag'
task_id = 'sleep_1000_secs'
with create_session() as session:
dag_bag = DagBag(
dag_folder=TEST_DAG_FOLDER,
include_examples=False,
)
dag = dag_bag.get_dag(dag_id)
task = dag.get_task(task_id)
dag.create_dagrun(
run_id="sleep_1000_secs_run",
state=State.RUNNING,
execution_date=DEFAULT_DATE,
start_date=DEFAULT_DATE,
session=session,
)
ti = TaskInstance(task=task, execution_date=DEFAULT_DATE)
ti.state = State.SCHEDULED
dag_model = DagModel(
dag_id=dag_id,
is_paused=False,
concurrency=5,
has_task_concurrency_limits=False,
)
session.merge(dag_model)
session.merge(ti)
session.commit()
executor = LocalExecutor(2)
executor.start()
executor.heartbeat()
executor.schedule_task(ti.key, SchedulingAction.START)
executor.heartbeat()
time.sleep(30)
# wait for task instance started
ti.refresh_from_db()
self.assertEqual(ti.state, State.RUNNING)
process = psutil.Process(ti.pid)
self.assertIsNotNone(process)
child = process.children(recursive=False)
self.assertEqual(1, len(child))
grandchild = child[0].children(recursive=False)
self.assertEqual(1, len(grandchild))
tes = self._check_task_execution(ti)
self.assertEqual(1, len(tes))
# restart the task instance
executor.schedule_task(ti.key, SchedulingAction.RESTART)
executor.heartbeat()
time.sleep(30)
self.assertFalse(self._check_process_exist(process.pid))
self.assertFalse(self._check_process_exist(child[0].pid))
self.assertFalse(self._check_process_exist(grandchild[0].pid))
ti.refresh_from_db()
self.assertEqual(ti.state, State.RUNNING)
process = psutil.Process(ti.pid)
self.assertIsNotNone(process)
child = process.children(recursive=False)
self.assertEqual(1, len(child))
grandchild = child[0].children(recursive=False)
self.assertEqual(1, len(grandchild))
tes = self._check_task_execution(ti)
self.assertEqual(2, len(tes))
self.assertEqual(2, tes[0].seq_num)
executor.schedule_task(ti.key, SchedulingAction.STOP)
ti.refresh_from_db()
time.sleep(10)
self.assertEqual(State.KILLED, ti.state)
self.assertFalse(self._check_process_exist(process.pid))
self.assertFalse(self._check_process_exist(child[0].pid))
self.assertFalse(self._check_process_exist(grandchild[0].pid))
self._check_task_execution(ti)
executor.end()