Merge pull request #1007 from mikegrima/db_fixes

Some database error fixes

security_monkey/task_scheduler/tasks.py

@@ -34,6 +34,13 @@ def task_account_tech(self, account_name, technology_name):
     app.logger.info("[ ] Executing Celery task for account: {}, technology: {}".format(account_name, technology_name))
     time1 = time.time()
 
+    # Verify that the account exists (was it deleted? was it renamed?):
+    if not Account.query.filter(Account.name == account_name).first():
+        app.logger.error("[X] Account has been removed or renamed: {}. Please restart the scheduler to fix.".format(
+            account_name
+        ))
+        return
+
     try:
         reporter_logic(account_name, technology_name)
 
@@ -57,6 +64,13 @@ def task_audit(self, account_name, technology_name):
 
     app.logger.info("[ ] Executing Celery task to audit changes for Account: {} Technology: {}".format(account_name,
                                                                                                        technology_name))
+    # Verify that the account exists (was it deleted? was it renamed?):
+    if not Account.query.filter(Account.name == account_name).first():
+        app.logger.error("[X] Account has been removed or renamed: {}. Please restart the scheduler to fix.".format(
+            account_name
+        ))
+        return
+
     try:
         audit_changes([account_name], [technology_name], True)
 

security_monkey/tests/scheduling/test_celery_scheduler.py

@@ -510,7 +510,6 @@ def test_acelery_skipabeat(self, mock_expired_exceptions, mock_account_tech, moc
         db.session.add(test_account)
         db.session.commit()
 
-
     @patch("security_monkey.task_scheduler.tasks.clear_old_exceptions")
     def test_celery_exception_task(self, mock_exception_clear):
         from security_monkey.task_scheduler.tasks import clear_expired_exceptions
@@ -558,3 +557,31 @@ def test_account_tech_task(self, mock_store_exception, mock_reporter, mock_setup
             task_account_tech(account_name, technology_name)  # noqa
 
         mock_store_exception.assert_called_with("scheduler-exception-on-watch", None, exception)
+
+    @patch("security_monkey.task_scheduler.tasks.setup")
+    @patch("security_monkey.task_scheduler.tasks.reporter_logic")
+    @patch("security_monkey.task_scheduler.tasks.app.logger.error")
+    def test_account_tech_no_account(self, mock_error, mock_reporter, mock_setup):
+        from security_monkey.task_scheduler.tasks import task_account_tech
+        technology_name = "iamrole"
+
+        task_account_tech("notanaccount", technology_name)  # noqa
+
+        assert not mock_reporter.called
+        assert mock_error.called
+        assert mock_error.call_args[0][0] == "[X] Account has been removed or renamed: notanaccount. " \
+                                             "Please restart the scheduler to fix."
+
+    @patch("security_monkey.task_scheduler.tasks.setup")
+    @patch("security_monkey.task_scheduler.tasks.audit_changes")
+    @patch("security_monkey.task_scheduler.tasks.app.logger.error")
+    def test_audit_no_account(self, mock_error, mock_audit_changes, mock_setup):
+        from security_monkey.task_scheduler.tasks import task_audit
+        technology_name = "iamrole"
+
+        task_audit("notanaccount", technology_name)  # noqa
+
+        assert not mock_audit_changes.called
+        assert mock_error.called
+        assert mock_error.call_args[0][0] == "[X] Account has been removed or renamed: notanaccount. " \
+                                             "Please restart the scheduler to fix."