From e2a5509b763272c0a538208fac019085eb7de1ca Mon Sep 17 00:00:00 2001 From: sadnub Date: Sat, 13 Feb 2021 14:38:03 -0500 Subject: [PATCH] add missing task tests in automation and alerts --- api/tacticalrmm/alerts/tasks.py | 5 - api/tacticalrmm/alerts/tests.py | 46 +++++- api/tacticalrmm/automation/tasks.py | 14 +- api/tacticalrmm/automation/tests.py | 220 +++------------------------- 4 files changed, 72 insertions(+), 213 deletions(-) diff --git a/api/tacticalrmm/alerts/tasks.py b/api/tacticalrmm/alerts/tasks.py index cc6ee2a3..cede7445 100644 --- a/api/tacticalrmm/alerts/tasks.py +++ b/api/tacticalrmm/alerts/tasks.py @@ -13,8 +13,3 @@ def unsnooze_alerts() -> str: ) return "ok" - - -@app.task -def periodic_alert_notifications() -> str: - return "not implemented" diff --git a/api/tacticalrmm/alerts/tests.py b/api/tacticalrmm/alerts/tests.py index e546a6c2..827f888a 100644 --- a/api/tacticalrmm/alerts/tests.py +++ b/api/tacticalrmm/alerts/tests.py @@ -1,6 +1,7 @@ from datetime import datetime, timedelta from core.models import CoreSettings +from django.utils import timezone as djangotime from tacticalrmm.test import TacticalTestCase from model_bakery import baker, seq @@ -10,7 +11,6 @@ from .serializers import ( AlertTemplateSerializer, AlertTemplateRelationSerializer, ) -from typing import List, Any class TestAlertsViews(TacticalTestCase): @@ -21,28 +21,38 @@ class TestAlertsViews(TacticalTestCase): def test_get_alerts(self): url = "/alerts/alerts/" + # create check, task, and agent to test each serializer function + check = baker.make_recipe("checks.diskspace_check") + task = baker.make("autotasks.AutomatedTask") + agent = baker.make_recipe("agents.agent") # setup data alerts = baker.make( "alerts.Alert", + agent=agent, alert_time=seq(datetime.now(), timedelta(days=15)), severity="warning", _quantity=3, ) baker.make( "alerts.Alert", + assigned_check=check, alert_time=seq(datetime.now(), timedelta(days=15)), severity="error", _quantity=7, ) baker.make( "alerts.Alert", + assigned_task=task, snoozed=True, + snooze_until=djangotime.now(), alert_time=seq(datetime.now(), timedelta(days=15)), _quantity=2, ) baker.make( "alerts.Alert", + agent=agent, resolved=True, + resolved_on=djangotime.now(), alert_time=seq(datetime.now(), timedelta(days=15)), _quantity=9, ) @@ -331,3 +341,37 @@ class TestAlertsViews(TacticalTestCase): self.assertTrue( AlertTemplate.objects.get(pk=alert_template.pk).is_default_template ) + + +class TestAlertTasks(TacticalTestCase): + def test_unsnooze_alert_task(self): + from alerts.tasks import unsnooze_alerts + + # these will be unsnoozed whent eh function is run + not_snoozed = baker.make( + "alerts.Alert", + snoozed=True, + snooze_until=seq(datetime.now(), timedelta(days=15)), + _quantity=5, + ) + + # these will still be snoozed after the function is run + snoozed = baker.make( + "alerts.Alert", + snoozed=True, + snooze_until=seq(datetime.now(), timedelta(days=-15)), + _quantity=5, + ) + + unsnooze_alerts() + + self.assertFalse( + Alert.objects.filter( + pk__in=[alert.pk for alert in not_snoozed], snoozed=False + ).exists() + ) + self.assertTrue( + Alert.objects.filter( + pk__in=[alert.pk for alert in snoozed], snoozed=False + ).exists() + ) diff --git a/api/tacticalrmm/automation/tasks.py b/api/tacticalrmm/automation/tasks.py index 0db0e986..60bc3cd3 100644 --- a/api/tacticalrmm/automation/tasks.py +++ b/api/tacticalrmm/automation/tasks.py @@ -7,6 +7,7 @@ from tacticalrmm.celery import app @app.task +# generates policy checks on agents affected by a policy and optionally generate automated tasks def generate_agent_checks_from_policies_task(policypk, create_tasks=False): policy = Policy.objects.get(pk=policypk) @@ -31,6 +32,7 @@ def generate_agent_checks_from_policies_task(policypk, create_tasks=False): @app.task +# generates policy checks on a list of agents and optionally generate automated tasks def generate_agent_checks_task(agentpks, create_tasks=False): for agent in Agent.objects.filter(pk__in=agentpks): agent.generate_checks_from_policies() @@ -40,6 +42,7 @@ def generate_agent_checks_task(agentpks, create_tasks=False): @app.task +# generates policy checks on agent servers or workstations within a certain client or site and optionally generate automated tasks def generate_agent_checks_by_location_task(location, mon_type, create_tasks=False): for agent in Agent.objects.filter(**location).filter(monitoring_type=mon_type): @@ -50,6 +53,7 @@ def generate_agent_checks_by_location_task(location, mon_type, create_tasks=Fals @app.task +# generates policy checks on all agent servers or workstations and optionally generate automated tasks def generate_all_agent_checks_task(mon_type, create_tasks=False): for agent in Agent.objects.filter(monitoring_type=mon_type): agent.generate_checks_from_policies() @@ -59,12 +63,14 @@ def generate_all_agent_checks_task(mon_type, create_tasks=False): @app.task +# deletes a policy managed check from all agents def delete_policy_check_task(checkpk): Check.objects.filter(parent_check=checkpk).delete() @app.task +# updates policy managed check fields on agents def update_policy_check_fields_task(checkpk): check = Check.objects.get(pk=checkpk) @@ -100,6 +106,7 @@ def update_policy_check_fields_task(checkpk): @app.task +# generates policy tasks on agents affected by a policy def generate_agent_tasks_from_policies_task(policypk): policy = Policy.objects.get(pk=policypk) @@ -121,13 +128,6 @@ def generate_agent_tasks_from_policies_task(policypk): agent.generate_tasks_from_policies() -@app.task -def generate_agent_tasks_by_location_task(location, mon_type): - - for agent in Agent.objects.filter(**location).filter(monitoring_type=mon_type): - agent.generate_tasks_from_policies() - - @app.task def delete_policy_autotask_task(taskpk): from autotasks.tasks import delete_win_task_schedule diff --git a/api/tacticalrmm/automation/tests.py b/api/tacticalrmm/automation/tests.py index 3080cbb4..d44c8e2f 100644 --- a/api/tacticalrmm/automation/tests.py +++ b/api/tacticalrmm/automation/tests.py @@ -907,206 +907,6 @@ class TestPolicyTasks(TacticalTestCase): self.assertEqual(task.parent_task, tasks[2].id) self.assertEqual(task.name, tasks[2].name) - @patch("automation.tasks.generate_agent_checks_by_location_task.delay") - def test_generate_agent_tasks_by_location( - self, generate_agent_checks_by_location_task - ): - from automation.tasks import ( - generate_agent_tasks_by_location_task as generate_agent_tasks, - ) - - # setup data - policy = baker.make("automation.Policy", active=True) - baker.make( - "autotasks.AutomatedTask", policy=policy, name=seq("Task"), _quantity=3 - ) - - server_agent = baker.make_recipe("agents.server_agent") - workstation_agent = baker.make_recipe("agents.workstation_agent") - - # nothing should have tasks yet - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 0 - ) - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 0) - - # set workstation policy on client and policy tasks should be there - workstation_agent.client.workstation_policy = policy - workstation_agent.client.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site__client_id": workstation_agent.client.pk}, - mon_type="workstation", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site__client_id": workstation_agent.client.pk}, - mon_type="workstation", - ) - - # make sure the tasks were added - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 0) - - # remove workstation policy from client - workstation_agent.client.workstation_policy = None - workstation_agent.client.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site__client_id": workstation_agent.client.pk}, - mon_type="workstation", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site__client_id": workstation_agent.client.pk}, - mon_type="workstation", - ) - - # make sure the tasks were removed - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 0) - - # set server policy on client and policy tasks should be there - server_agent.client.server_policy = policy - server_agent.client.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site__client_id": server_agent.client.pk}, - mon_type="server", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site__client_id": server_agent.client.pk}, mon_type="server" - ) - - # make sure the tasks were added (Will be there but pending deletion) - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3) - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - - # remove server policy from client - server_agent.client.server_policy = None - server_agent.client.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site__client_id": server_agent.client.pk}, - mon_type="server", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site__client_id": server_agent.client.pk}, mon_type="server" - ) - - # make sure the tasks were removed - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3) - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - - # set workstation policy on site and policy checks should be there - workstation_agent.site.workstation_policy = policy - workstation_agent.site.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site_id": workstation_agent.site.pk}, - mon_type="workstation", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site_id": workstation_agent.client.pk}, mon_type="workstation" - ) - - # make sure the tasks were added - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3) - - # remove workstation policy from site - workstation_agent.site.workstation_policy = None - workstation_agent.site.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site_id": workstation_agent.site.pk}, - mon_type="workstation", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site_id": workstation_agent.client.pk}, mon_type="workstation" - ) - - # make sure tasks were removed - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3) - - # set server policy on site and policy checks should be there - server_agent.site.server_policy = policy - server_agent.site.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site_id": server_agent.site.pk}, - mon_type="server", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site_id": server_agent.client.pk}, mon_type="server" - ) - - # make sure tasks were added - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3) - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - - # remove server policy from site - server_agent.site.server_policy = None - server_agent.site.save() - - # should trigger task in save method on core - generate_agent_checks_by_location_task.assert_called_with( - location={"site_id": server_agent.site.pk}, - mon_type="server", - create_tasks=True, - ) - generate_agent_checks_by_location_task.reset_mock() - - generate_agent_tasks( - location={"site_id": server_agent.client.pk}, mon_type="server" - ) - - self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 3) - self.assertEqual( - Agent.objects.get(pk=workstation_agent.id).autotasks.count(), 3 - ) - @patch("autotasks.tasks.delete_win_task_schedule.delay") def test_delete_policy_tasks(self, delete_win_task_schedule): from .tasks import delete_policy_autotask_task @@ -1157,3 +957,23 @@ class TestPolicyTasks(TacticalTestCase): enable_or_disable_win_task.assert_called_with( agent.autotasks.get(parent_task=tasks[0].id).id, False ) + + @patch("agents.models.Agent.generate_tasks_from_policies") + @patch("agents.models.Agent.generate_checks_from_policies") + def test_generate_agent_checks_with_agentpks(self, generate_checks, generate_tasks): + from automation.tasks import generate_agent_checks_task + + agents = baker.make_recipe("agents.agent", _quantity=5) + + # reset because creating agents triggers it + generate_checks.reset_mock() + generate_tasks.reset_mock() + + generate_agent_checks_task([agent.pk for agent in agents]) + self.assertEquals(generate_checks.call_count, 5) + generate_tasks.assert_not_called() + generate_checks.reset_mock() + + generate_agent_checks_task([agent.pk for agent in agents], create_tasks=True) + self.assertEquals(generate_checks.call_count, 5) + self.assertEquals(generate_checks.call_count, 5)