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