131 lines
8.1 KiB
Python
131 lines
8.1 KiB
Python
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from unittest.mock import patch
|
|
import task_core as core
|
|
import taskd
|
|
|
|
class TaskTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp=tempfile.TemporaryDirectory(prefix='scheduled-task-qa-')
|
|
self.original=(core.BASE,core.LOGS,core.UNITS,taskd.ADMIN)
|
|
folder=Path(self.temp.name)
|
|
core.BASE=folder/'data';core.LOGS=folder/'logs';core.UNITS=folder/'units';core.UNITS.mkdir()
|
|
taskd.ADMIN=folder/'admin.json';taskd.SESSIONS.clear();taskd.ATTEMPTS.clear();core.init_store()
|
|
core.atom_write(taskd.ADMIN,json.dumps(taskd.password_record('only-a-test-password')))
|
|
|
|
def tearDown(self):
|
|
core.BASE,core.LOGS,core.UNITS,taskd.ADMIN=self.original
|
|
taskd.SESSIONS.clear();taskd.ATTEMPTS.clear();self.temp.cleanup()
|
|
|
|
def definition(self,**values):
|
|
return {'name':'测试任务','mode':'command','body':'printf hello','enabled':False,
|
|
'schedule':{'kind':'daily','time':'03:00'},**values}
|
|
|
|
def test_daily_weekly_monthly_and_interval(self):
|
|
self.assertEqual(core.schedule_definition({'kind':'daily','time':'03:15'})[0],['OnCalendar=*-*-* 03:15:00 Asia/Shanghai'])
|
|
self.assertIn('Mon,Sun',core.schedule_definition({'kind':'weekly','weekdays':[0,6],'time':'04:00'})[0][0])
|
|
self.assertIn('跳过',core.schedule_definition({'kind':'monthly','day':31,'time':'02:00'})[1])
|
|
interval=core.schedule_definition({'kind':'minutes','every':10})[0]
|
|
self.assertEqual(interval,['OnActiveSec=600s','OnUnitInactiveSec=600s'])
|
|
self.assertFalse(any('OnBootSec' in x for x in interval))
|
|
|
|
def test_cron_lists_ranges_steps_and_day_or(self):
|
|
output=core.cron_calendars('*/15 3 * * MON-FRI')
|
|
self.assertEqual(len(output),1)
|
|
self.assertIn('0,15,30,45:00',output[0])
|
|
self.assertIn('Mon,Tue,Wed,Thu,Fri',output[0])
|
|
output=core.cron_calendars('0 3 1 * 1')
|
|
self.assertEqual(len(output),2)
|
|
self.assertIn('*-*-1',output[0]);self.assertIn('Mon',output[1])
|
|
for bad in ['0 3 * *','*/0 * * * *','0 24 * * *','0 3 32 * *','0 3 * * 9','0 3 L * *']:
|
|
with self.subTest(bad=bad),self.assertRaises(core.TaskError):core.cron_calendars(bad)
|
|
|
|
def test_task_text_never_becomes_unit_instructions(self):
|
|
definition=core.validate_definition(self.definition(body='echo test\n# ExecStart=/bad/example'),check_paths=False)
|
|
service,timer=core.render_units('a'*32,definition)
|
|
self.assertNotIn('echo test',service);self.assertNotIn('/bad/example',service)
|
|
self.assertIn('task_exec.py '+'a'*32,service)
|
|
self.assertIn('Persistent=false',timer)
|
|
with self.assertRaises(core.TaskError):core.validate_definition(self.definition(name='name\nExecStart=bad'),check_paths=False)
|
|
with self.assertRaises(core.TaskError):core.validate_id('../../etc/shadow')
|
|
|
|
def test_validation_limits_and_past_dates(self):
|
|
for changes in [{'body':''},{'timeout':1},{'timeout':90000},{'body':'x'*65537},{'enabled':'false'},
|
|
{'schedule':{'kind':'once','at':'2020-01-01T03:00'}}]:
|
|
with self.subTest(changes=list(changes)),self.assertRaises(core.TaskError):core.validate_definition(self.definition(**changes),check_paths=False)
|
|
self.assertTrue(core.validate_definition(self.definition(body='systemctl reboot'),check_paths=False)['dangerous'])
|
|
|
|
def test_session_and_csrf_gate_privileged_operations(self):
|
|
with patch.object(core,'save_task') as save:
|
|
with self.assertRaises(core.TaskError) as error:taskd.dispatch({'action':'tasks.save','payload':self.definition()})
|
|
self.assertEqual(error.exception.status,401);save.assert_not_called()
|
|
login=taskd.login({'username':'admin','password':'only-a-test-password'},'qa')
|
|
with patch.object(core,'set_enabled') as toggle:
|
|
with self.assertRaises(core.TaskError) as error:taskd.dispatch({'action':'tasks.enable','token':login['token'],'csrf':'wrong','payload':{'id':'a'*32,'enabled':True}})
|
|
self.assertEqual(error.exception.status,403);toggle.assert_not_called()
|
|
self.assertTrue(taskd.dispatch({'action':'auth.status','token':login['token']})['logged_in'])
|
|
taskd.dispatch({'action':'auth.logout','token':login['token'],'csrf':login['csrf']})
|
|
self.assertFalse(taskd.dispatch({'action':'auth.status','token':login['token']})['logged_in'])
|
|
|
|
def test_password_change_revokes_all_sessions(self):
|
|
one=taskd.login({'username':'admin','password':'only-a-test-password'},'qa')
|
|
two=taskd.login({'username':'admin','password':'only-a-test-password'},'qa-2')
|
|
taskd.dispatch({'action':'auth.password','token':one['token'],'csrf':one['csrf'],
|
|
'payload':{'old_password':'only-a-test-password','new_password':'new-test-password-123'}})
|
|
self.assertFalse(taskd.verify('only-a-test-password'));self.assertTrue(taskd.verify('new-test-password-123'))
|
|
self.assertFalse(taskd.dispatch({'action':'auth.status','token':two['token']})['logged_in'])
|
|
|
|
def test_wrong_old_password_keeps_current_login(self):
|
|
login=taskd.login({'username':'admin','password':'only-a-test-password'},'qa')
|
|
with self.assertRaises(core.TaskError) as error:
|
|
taskd.dispatch({'action':'auth.password','token':login['token'],'csrf':login['csrf'],
|
|
'payload':{'old_password':'wrong-password','new_password':'new-test-password-123'}})
|
|
self.assertEqual(error.exception.status,400)
|
|
self.assertTrue(taskd.dispatch({'action':'auth.status','token':login['token']})['logged_in'])
|
|
|
|
def test_log_tail_is_bounded_and_missing_id_rejected(self):
|
|
task_id='b'*32;run_id='c'*32
|
|
with core.db() as con:
|
|
con.execute('INSERT INTO runs(id,task_id,task_name,started,state) VALUES(?,?,?,?,?)',(run_id,task_id,'test',time.time(),'success'))
|
|
(core.LOGS/(run_id+'.log')).write_text('x'*200000)
|
|
result=core.read_run(run_id)
|
|
self.assertTrue(result['tail_only']);self.assertEqual(len(result['output']),131072)
|
|
with self.assertRaises(core.TaskError):core.read_run('../password')
|
|
|
|
def test_gpu_count_requires_login_csrf_and_preserves_other_config(self):
|
|
config=Path(self.temp.name)/'hardware-watch.json'
|
|
config.write_text(json.dumps({'expected_gpu_count':2,'expected_gpu_uuids':['old-p100'],
|
|
'expected_disks':['nvme0n1'],'gpu_warn_c':85}))
|
|
with patch.object(taskd,'MONITOR_CONFIG',config),patch.object(core,'execute') as execute:
|
|
execute.return_value.returncode=0
|
|
request={'action':'monitor.gpu_count','payload':{'expected_gpu_count':1}}
|
|
with self.assertRaises(core.TaskError) as error:taskd.dispatch(request)
|
|
self.assertEqual(error.exception.status,401)
|
|
login=taskd.login({'username':'admin','password':'only-a-test-password'},'qa')
|
|
request['token']=login['token']
|
|
with self.assertRaises(core.TaskError) as error:taskd.dispatch(request)
|
|
self.assertEqual(error.exception.status,403)
|
|
request['csrf']=login['csrf']
|
|
for bad in [True,0,3,'1',1.0]:
|
|
request['payload']={'expected_gpu_count':bad}
|
|
with self.subTest(value=bad),self.assertRaises(core.TaskError):taskd.dispatch(request)
|
|
for count in (1,2):
|
|
request['payload']={'expected_gpu_count':count}
|
|
taskd.dispatch(request)
|
|
saved=json.loads(config.read_text())
|
|
self.assertEqual(saved['expected_gpu_count'],count)
|
|
self.assertEqual(saved['expected_gpu_uuids'],[])
|
|
self.assertEqual(saved['expected_disks'],['nvme0n1'])
|
|
self.assertEqual(saved['gpu_warn_c'],85)
|
|
self.assertEqual(execute.call_count,2)
|
|
|
|
def test_night_cycle_is_identified_as_shutdown_task(self):
|
|
definition=core.validate_definition(self.definition(body='/usr/bin/python3 /opt/hardware-dashboard/scripts/night-power-cycle.py'),check_paths=False)
|
|
self.assertTrue(definition['dangerous'])
|
|
|
|
if __name__=='__main__':unittest.main(verbosity=2)
|