def test_on_kill(self): """ Test that ensures that clearing in the UI SIGTERMS the task """ path = "/tmp/airflow_on_kill" try: os.unlink(path) except OSError: pass dagbag = models.DagBag( dag_folder=TEST_DAG_FOLDER, include_examples=False, ) dag = dagbag.dags.get('test_on_kill') task = dag.get_task('task1') session = settings.Session() dag.clear() dag.create_dagrun(run_id="test", state=State.RUNNING, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, session=session) ti = TI(task=task, execution_date=DEFAULT_DATE) job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True) runner = StandardTaskRunner(job1) runner.start() # give the task some time to startup time.sleep(3) runner.terminate() f = open(path, "r") self.assertEqual("ON_KILL_TEST", f.readline()) f.close()
def test_start_and_terminate(self): local_task_job = mock.Mock() local_task_job.task_instance = mock.MagicMock() local_task_job.task_instance.run_as_user = None local_task_job.task_instance.command_as_list.return_value = ['sleep', '1000'] runner = StandardTaskRunner(local_task_job) runner.start() pgid = os.getpgid(runner.process.pid) self.assertTrue(pgid) procs = [] for p in psutil.process_iter(): try: if os.getpgid(p.pid) == pgid: procs.append(p) except OSError: pass runner.terminate() for p in procs: self.assertFalse(psutil.pid_exists(p.pid))
def test_on_kill(self): """ Test that ensures that clearing in the UI SIGTERMS the task """ path = "/tmp/airflow_on_kill" try: os.unlink(path) except OSError: pass dagbag = models.DagBag( dag_folder=TEST_DAG_FOLDER, include_examples=False, ) dag = dagbag.dags.get('test_on_kill') task = dag.get_task('task1') session = settings.Session() dag.clear() dag.create_dagrun(run_id="test", state=State.RUNNING, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE, session=session) ti = TI(task=task, execution_date=DEFAULT_DATE) job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True) runner = StandardTaskRunner(job1) runner.start() # Give the task some time to startup time.sleep(10) runner.terminate() # Wait some time for the result for _ in range(20): if os.path.exists(path): break time.sleep(2) with open(path, "r") as f: self.assertEqual("ON_KILL_TEST", f.readline())