def successful_remediation(remediation: database.Remediation) -> bool:
"""Handle successful remediations"""
if remediation.result not in [MailOutcome.REMOVED, MailOutcome.RESTORED]:
return False
remediation.successful = True
remediation.status = REMEDIATION_STATUS_COMPLETED
logging.info(f'completed remediation request {remediation}')
return True
def request(action,
type,
key,
user_id,
company_id,
comment=None,
lock=None,
lock_time=None,
status=REMEDIATION_STATUS_NEW):
remediation = Remediation(
action=action,
type=type,
key=key,
user_id=user_id,
comment=comment,
company_id=company_id,
lock=lock,
lock_time=lock_time,
status=status,
)
saq.db.add(remediation)
saq.db.commit()
saq.db.refresh(remediation)
saq.db.expunge_all()
return remediation
def test_successful_remediation_restored_success(self):
class Remediation:
def __init__(self):
self.result = 'restored'
self.successful = None
self.status = None
r = Remediation()
result = mail.successful_remediation(r)
self.assertTrue(result)
self.assertTrue(r.successful)
self.assertEqual('COMPLETED', r.status)
def test_successful_remediation_failure(self):
class Remediation:
def __init__(self):
self.result = 'invalid'
self.successful = None
self.status = None
r = Remediation()
result = mail.successful_remediation(r)
self.assertFalse(result)
self.assertIsNone(r.successful)
self.assertIsNone(r.status)
def test_failed_remediation_mail_outcome_other_error(self):
class Remediation:
def __init__(self):
self.result = 'non-standard-error'
r = Remediation()
class Remediator:
def __init__(self):
self.config_name = 'my_config'
r2 = Remediator()
error_dict = {'my_config': None}
mail.failed_remediation(r, r2, error_dict)
self.assertEqual('non-standard-error', error_dict['my_config'])
def test_failed_remediation_mail_outcome_error(self):
class Remediation:
def __init__(self):
self.result = 'error'
r = Remediation()
class Remediator:
def __init__(self):
self.config_name = 'my_config'
r2 = Remediator()
error_dict = {'my_config': None}
mail.failed_remediation(r, r2, error_dict)
self.assertEqual('unknown error while remediating', error_dict['my_config'])
def test_attempt_remediation_return_true(self):
class Remediation:
def __init__(self):
self.result = None
self.action = 'remove'
class Remediator:
def __init__(self):
self.config_name = 'my_config'
def remediate(self, *args, **kwargs):
return 'expected'
r = Remediation()
attempt = mail.attempt_remediation(r, Remediator(), '*****@*****.**', '<*****@*****.**>', {})
self.assertTrue(attempt)
self.assertEqual(r.result, 'expected')
def attempt_remediation(remediation: database.Remediation,
remediator: BaseRemediator, recipient: str,
message_id: str, error_dict: dict) -> bool:
"""Helper function to handle attempting remediation attempt and any unhandled errors."""
try:
remediation.result = remediator.remediate(remediation.action,
recipient, message_id)
except Exception as e:
# This should be caught within remediate()... but just in case the
# developer didn't implement error handling correctly...
logging.error(f"uncaught error when remediating: {e.__class__} '{e}'")
error_dict[remediator.config_name] = 'uncaught error while remediating'
return False
else:
return True
def test_attempt_remediation_return_false(self):
class Remediation:
def __init__(self):
self.result = None
self.action = 'remove'
class Remediator:
def __init__(self):
self.config_name = 'my_config'
def remediate(self, *args, **kwargs):
raise ValueError('expect me')
r = Remediation()
error_dict = {'my_config': None}
attempt = mail.attempt_remediation(r, Remediator(), '*****@*****.**', '<*****@*****.**>', error_dict)
self.assertFalse(attempt)
self.assertIsNone(r.result)
self.assertEqual('uncaught error while remediating', error_dict['my_config'])
def _insert_email_remediation_object(action,
message_id,
recipient,
user_id,
company_id,
comment=None):
remediation = Remediation(type=REMEDIATION_TYPE_EMAIL,
action=action,
user_id=user_id,
key=f'{message_id}:{recipient}',
comment=comment,
company_id=company_id)
saq.db.add(remediation)
saq.db.commit()
logging.info(
f"user {user_id} added remediation request for message_id {message_id} recipient {recipient}"
)
return True