def test_kill_zombies_when_job_state_is_not_running(self, mock_ti_handle_failure): """ Test that kill zombies calls TI's failure handler with proper context """ dagbag = models.DagBag(dag_folder=self.empty_dir, include_examples=True) with create_session() as session: session.query(TI).delete() session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() dagbag.kill_zombies() mock_ti_handle_failure \ .assert_called_with(ANY, configuration.getboolean('core', 'unit_test_mode'), ANY)
def test_kill_zombie_when_job_received_no_heartbeat(self, mock_ti_handle_failure): """ Test that kill zombies calls TI's failure handler with proper context """ zombie_threshold_secs = ( configuration.getint('scheduler', 'scheduler_zombie_task_threshold')) dagbag = models.DagBag(dag_folder=self.empty_dir, include_examples=True) with create_session() as session: session.query(TI).delete() session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.latest_heartbeat = utcnow() - timedelta(seconds=zombie_threshold_secs) lj.state = State.RUNNING lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() dagbag.kill_zombies() mock_ti_handle_failure \ .assert_called_with(ANY, configuration.getboolean('core', 'unit_test_mode'), ANY)
def test_kill_zombies_doesn_nothing(self, mock_ti_handle_failure): """ Test that kill zombies does nothing when job is running and received heartbeat """ dagbag = models.DagBag(dag_folder=self.empty_dir, include_examples=True) with create_session() as session: session.query(TI).delete() session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.latest_heartbeat = utcnow() lj.state = State.RUNNING lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() dagbag.kill_zombies() mock_ti_handle_failure.assert_not_called()
def test_zombies_are_correctly_passed_to_dag_file_processor(self): """ Check that the same set of zombies are passed to the dag file processors until the next zombie detection logic is invoked. """ test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py') with conf_vars({ ('scheduler', 'max_threads'): '1', ('core', 'load_examples'): 'False' }): dagbag = DagBag(test_dag_path) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('test_example_bash_operator') dag.sync_to_db() task = dag.get_task(task_id='run_this_last') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() fake_zombies = [SimpleTaskInstance(ti)] async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn') processor_agent = DagFileProcessorAgent( test_dag_path, [], 1, FakeDagFileProcessorRunner._fake_dag_processor_factory, timedelta.max, [], False, async_mode) processor_agent.start() parsing_result = [] if not async_mode: processor_agent.heartbeat() while not processor_agent.done: if not async_mode: processor_agent.wait_until_finished() parsing_result.extend(processor_agent.harvest_simple_dags()) self.assertEqual(len(fake_zombies), len(parsing_result)) self.assertEqual(set([zombie.key for zombie in fake_zombies]), set([result.key for result in parsing_result]))
def test_find_zombies(self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True) dagbag = DagBag(TEST_DAG_FOLDER) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() manager._last_zombie_query_time = timezone.utcnow() - timedelta( seconds=manager._zombie_threshold_secs + 1) manager._find_zombies() zombies = manager._zombies self.assertEqual(1, len(zombies)) self.assertIsInstance(zombies[0], SimpleTaskInstance) self.assertEqual(ti.dag_id, zombies[0].dag_id) self.assertEqual(ti.task_id, zombies[0].task_id) self.assertEqual(ti.execution_date, zombies[0].execution_date) session.query(TI).delete() session.query(LJ).delete()
def test_find_zombies(self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, signal_conn=MagicMock(), stat_queue=MagicMock(), result_queue=MagicMock, async_mode=True) dagbag = DagBag(TEST_DAG_FOLDER) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() manager._last_zombie_query_time = timezone.utcnow() - timedelta( seconds=manager._zombie_threshold_secs + 1) zombies = manager._find_zombies() self.assertEqual(1, len(zombies)) self.assertIsInstance(zombies[0], SimpleTaskInstance) self.assertEqual(ti.dag_id, zombies[0].dag_id) self.assertEqual(ti.task_id, zombies[0].task_id) self.assertEqual(ti.execution_date, zombies[0].execution_date) session.query(TI).delete() session.query(LJ).delete()
def test_zombies_are_correctly_passed_to_dag_file_processor(self): """ Check that the same set of zombies are passed to the dag file processors until the next zombie detection logic is invoked. """ with conf_vars({('scheduler', 'max_threads'): '1', ('core', 'load_examples'): 'False'}): dagbag = DagBag(os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py')) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('test_example_bash_operator') task = dag.get_task(task_id='run_this_last') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() fake_zombies = [SimpleTaskInstance(ti)] class FakeDagFIleProcessor(DagFileProcessor): # This fake processor will return the zombies it received in constructor # as its processing result w/o actually parsing anything. def __init__(self, file_path, pickle_dags, dag_id_white_list, zombies): super(FakeDagFIleProcessor, self).__init__( file_path, pickle_dags, dag_id_white_list, zombies ) self._result = zombies, 0 def start(self): pass @property def start_time(self): return DEFAULT_DATE @property def pid(self): return 1234 @property def done(self): return True @property def result(self): return self._result def processor_factory(file_path, zombies): return FakeDagFIleProcessor(file_path, False, [], zombies) test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py') async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn') processor_agent = DagFileProcessorAgent(test_dag_path, [], 1, processor_factory, timedelta.max, async_mode) processor_agent.start() parsing_result = [] if not async_mode: processor_agent.heartbeat() while not processor_agent.done: if not async_mode: processor_agent.wait_until_finished() parsing_result.extend(processor_agent.harvest_simple_dags()) self.assertEqual(len(fake_zombies), len(parsing_result)) self.assertEqual(set([zombie.key for zombie in fake_zombies]), set([result.key for result in parsing_result]))