def test_get_scan(self):
scan_id = "scan_id"
Scan.get_scan(scan_id, self.user_id, self.session)
self.session.query.assert_called_once_with(Scan)
m_filter = self.session.query(Scan).filter
m_filter.assert_called_once()
m_filter().one.assert_called_once()
def new(frontend_scan_id, user, session):
try:
scan = Scan.get_scan(frontend_scan_id, user.id, session)
except IrmaDatabaseResultNotFound:
scan = Scan(frontend_scan_id, user.id)
session.add(scan)
session.commit()
log.debug("scanid %s: user_id %s id %s", frontend_scan_id, user.id,
scan.id)
return scan
def new(frontend_scan_id, user, nb_files, session):
try:
scan = Scan.get_scan(frontend_scan_id, user.id, session)
scan.nb_files += nb_files
scan.update(['nb_files'], session)
except IrmaDatabaseResultNotFound:
scan = Scan(frontend_scan_id, user.id, nb_files)
scan.save(session)
session.commit()
log.debug("scanid %s: user_id %s nb_files %s id %s", frontend_scan_id,
user.id, nb_files, scan.id)
return scan
def test_nb_files(self):
with patch('brain.models.sqlobjects.Scan.files',
new_callable=PropertyMock) as mock_files:
mock_files.return_value = {}
scan = Scan(self.frontend_scanid, self.user_id)
self.assertEqual(scan.nb_files, 0)
mock_files.assert_called_once()
def test_scan_cancel(self):
scanid = scan_ctrl.new(self.scanid, self.userid, 10)
for code in IrmaScanStatus.label.keys():
scan_ctrl.error(scanid, code)
with session_query() as session:
scan = Scan.load(scanid, session)
self.assertEqual(scan.status, code)
def test001___init__(self):
scan = Scan(self.frontend_scanid, self.user_id, self.nb_files)
self.assertEqual(scan.scan_id, self.frontend_scanid)
self.assertEqual(scan.user_id, self.user_id)
self.assertEqual(scan.nb_files, self.nb_files)
self.assertEqual(scan.status, IrmaScanStatus.empty)
self.assertIsNotNone(scan.timestamp)
def flush(scan_id):
with session_transaction() as session:
scan = Scan.load(scan_id, session)
if scan.status == IrmaScanStatus.flushed:
return
for job in scan.jobs:
session.delete(job)
scan.status = IrmaScanStatus.flushed
def test_scan_new_status(self):
# test we have corrects fields
scanid = scan_ctrl.new(self.scanid, self.userid, 10)
with session_query() as session:
scan = Scan.load(scanid, session)
self.assertIsNotNone(scan.timestamp)
self.assertEqual(scan.status, IrmaScanStatus.empty)
self.assertEqual(scan.nb_files, 10)
self.assertIsNotNone(scan.nb_jobs)
def check_finished(scan_id):
with session_transaction() as session:
scan = Scan.load(scan_id, session)
if scan.status == IrmaScanStatus.processed:
return True
if scan.finished():
scan.status = IrmaScanStatus.processed
return True
return False
def scan_flush(scan_id):
try:
log.info("scan %s: flush requested", scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
scan_ctrl.flush(scan, session)
except Exception as e:
log.exception(type(e).__name__ + " : " + str(e))
return
def scan_flush(scan_id):
try:
log.info("scan %s: flush requested", scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
scan_ctrl.flush(scan, session)
except Exception as e:
log.exception(e)
return
def new(frontend_scan_id, user, session):
try:
scan = Scan.get_scan(frontend_scan_id, user.id, session)
except IrmaDatabaseResultNotFound:
scan = Scan(frontend_scan_id, user.id)
session.add(scan)
session.commit()
log.debug("scanid %s: user_id %s id %s",
frontend_scan_id, user.id, scan.id)
return scan
def scan_cancel(frontend_scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(frontend_scan_id, user.id, session)
res = scan_ctrl.cancel(scan, session)
return IrmaTaskReturn.success(res)
except Exception as e:
log.exception(e)
return IrmaTaskReturn.error("cancel error on brain")
def scan_cancel(frontend_scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(frontend_scan_id, user.id, session)
res = scan_ctrl.cancel(scan, session)
return IrmaTaskReturn.success(res)
except Exception as e:
log.exception(e)
return IrmaTaskReturn.error("cancel error on brain")
def scan_flush(frontend_scan_id):
try:
log.info("scan_id %s: scan flush requested", frontend_scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(frontend_scan_id, user.id, session)
scan_ctrl.flush(scan, session)
except Exception as e:
log.exception(e)
return
def files_flush(files, scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
ftpuser = scan.user.ftpuser
log.debug("Flushing files %s", files)
ftp_ctrl.flush(ftpuser, files)
except Exception as e:
log.exception(type(e).__name__ + " : " + str(e))
return
def test_scan_launched(self):
scan_id = scan_ctrl.new(self.scanid, self.userid, 10)
for i in xrange(0, 10):
for probe in ['probe1', 'probe2']:
job_ctrl.new(scan_id, "file-{0}".format(i), probe, 1)
scan_ctrl.launched(scan_id)
with session_query() as session:
scan = Scan.load(scan_id, session)
self.assertEqual(scan.status, IrmaScanStatus.launched)
self.assertEqual(scan.nb_files, 10)
self.assertEqual(scan.nb_jobs, 20)
def scan_cancel(scan_id):
try:
log.info("scan %s: cancel", scan_id)
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
res = scan_ctrl.cancel(scan, session)
return IrmaTaskReturn.success(res)
except Exception as e:
log.exception(type(e).__name__ + " : " + str(e))
return IrmaTaskReturn.error("cancel error on brain")
def files_flush(files, scan_id):
try:
with session_transaction() as session:
user = User.get_by_rmqvhost(session)
scan = Scan.get_scan(scan_id, user.id, session)
ftpuser = scan.user.ftpuser
log.debug("Flushing files %s", files)
ftp_ctrl.flush(ftpuser, files)
except Exception as e:
log.exception(e)
return
def setUp(self):
self.session = MagicMock()
self.name = "user"
self.rmqvhost = "vhost"
self.ftpuser = "******"
self.user = User(self.name, self.rmqvhost, self.ftpuser)
self.frontend_scanid = "frontend_scanid"
self.nb_files = randint(50, 100)
self.scan = Scan(frontend_scanid=self.frontend_scanid,
user_id=self.user.id)
self.scan.user = self.user
def progress(scan_id):
with session_query() as session:
scan = Scan.load(scan_id, session)
if IrmaScanStatus.is_error(scan.status):
status_str = IrmaScanStatus.label[scan.status]
raise IrmaValueError(status_str)
status = IrmaScanStatus.label[scan.status]
progress_details = None
if scan.status == IrmaScanStatus.launched:
(total, finished, success) = scan.progress()
progress_details = {}
progress_details['total'] = total
progress_details['finished'] = finished
progress_details['successful'] = success
return (status, progress_details)
def test_scan_job_error(self):
scan_id = scan_ctrl.new(self.scanid, self.userid, 10)
job_ids = []
for i in xrange(0, 10):
for probe in ['probe1', 'probe2']:
job_ids.append(job_ctrl.new(scan_id, "file-{0}".format(i),
probe, 1))
scan_ctrl.launched(scan_id)
shuffle(job_ids)
for job_id in job_ids:
job_ctrl.error(job_id)
self.assertTrue(scan_ctrl.check_finished(scan_id))
with session_query() as session:
scan = Scan.load(scan_id, session)
self.assertEqual(scan.status, IrmaScanStatus.processed)
self.assertEqual(scan.nb_files, 10)
self.assertEqual(scan.nb_jobs, 20)
def test_get_scan_multiple_found(self):
self.session.query.side_effect = MultipleResultsFound
with self.assertRaises(IrmaDatabaseError):
Scan.get_scan("whatever", self.user_id, self.session)
def test_get_scan_not_found(self):
self.session.query.side_effect = NoResultFound
with self.assertRaises(IrmaDatabaseResultNotFound):
Scan.get_scan("whatever", self.user_id, self.session)
def error(scan_id, code):
with session_transaction() as session:
scan = Scan.load(scan_id, session)
scan.status = code
def new(frontend_scan_id, user_id, nb_files):
with session_transaction() as session:
scan = Scan(frontend_scan_id, user_id, nb_files)
scan.save(session)
session.commit()
return scan.id
def get_scan_id(frontend_scan_id, user_id):
with session_query() as session:
scan = Scan.get_scan(frontend_scan_id, user_id, session)
return scan.id
def get_nbjobs(scan_id):
with session_query() as session:
scan = Scan.load(scan_id, session)
return scan.nb_jobs
def _set_status(scan_id, code):
with session_transaction() as session:
scan = Scan.load(scan_id, session)
scan.status = code
session.commit()
def get_pending_jobs(scan_id):
with session_query() as session:
scan = Scan.load(scan_id, session)
return scan.get_pending_jobs_taskid()
