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()