fix tests and added soem minor fixes

.devcontainer/docker-compose.yml

@@ -3,6 +3,7 @@ version: '3.4'
 services:
   api-dev:
     image: api-dev
+    restart: always
     build:
       context: .
       dockerfile: ./api.dockerfile
@@ -21,6 +22,7 @@ services:
 
   app-dev:
     image: node:12-alpine
+    restart: always
     command: /bin/sh -c "npm install && npm run serve -- --host 0.0.0.0 --port ${APP_PORT}"
     working_dir: /workspace/web
     volumes:

api/tacticalrmm/agents/baker_recipes.py

@@ -3,19 +3,20 @@ import string
 import os
 import json
 
-from model_bakery.recipe import Recipe, seq
+from model_bakery.recipe import Recipe, foreign_key
 from itertools import cycle
 from django.utils import timezone as djangotime
 from django.conf import settings
 
-from .models import Agent
-
 
 def generate_agent_id(hostname):
     rand = "".join(random.choice(string.ascii_letters) for _ in range(35))
     return f"{rand}-{hostname}"
 
 
+site = Recipe("clients.Site")
+
+
 def get_wmi_data():
     with open(
         os.path.join(settings.BASE_DIR, "tacticalrmm/test_data/wmi_python_agent.json")
@@ -24,7 +25,8 @@ def get_wmi_data():
 
 
 agent = Recipe(
-    Agent,
+    "agents.Agent",
+    site=foreign_key(site),
     hostname="DESKTOP-TEST123",
     version="1.3.0",
     monitoring_type=cycle(["workstation", "server"]),

api/tacticalrmm/agents/tests.py

@@ -705,6 +705,7 @@ class TestAgentViews(TacticalTestCase):
 class TestAgentViewsNew(TacticalTestCase):
     def setUp(self):
         self.authenticate()
+        self.setup_coresettings()
 
     def test_agent_counts(self):
         url = "/agents/agent_counts/"

api/tacticalrmm/automation/serializers.py

@@ -1,3 +1,4 @@
+from django.db.models.base import Model
 from rest_framework.serializers import (
     ModelSerializer,
     SerializerMethodField,

api/tacticalrmm/automation/tasks.py

@@ -146,7 +146,8 @@ def run_win_policy_autotask_task(task_pks):
 
 
 @app.task
-def update_policy_task_fields_task(taskpk):
+def update_policy_task_fields_task(taskpk, update_agent=False):
+    from autotasks.tasks import enable_or_disable_win_task
 
     task = AutomatedTask.objects.get(pk=taskpk)
 
@@ -160,4 +161,8 @@ def update_policy_task_fields_task(taskpk):
         name=task.name,
         timeout=task.timeout,
         enabled=task.enabled,
-    )
+    )
+
+    if update_agent:
+        for task in AutomatedTask.objects.filter(parent_task=taskpk):
+            enable_or_disable_win_task.delay(task.pk, task.enabled)

api/tacticalrmm/automation/tests.py

@@ -88,7 +88,7 @@ class TestPolicyViews(TacticalTestCase):
         self.check_not_authenticated("post", url)
 
     @patch("automation.tasks.generate_agent_checks_from_policies_task.delay")
-    def test_update_policy(self, mock_checks_task):
+    def test_update_policy(self, generate_agent_checks_from_policies_task):
         # returns 404 for invalid policy pk
         resp = self.client.put("/automation/policies/500/", format="json")
         self.assertEqual(resp.status_code, 404)
@@ -107,7 +107,7 @@ class TestPolicyViews(TacticalTestCase):
         self.assertEqual(resp.status_code, 200)
 
         # only called if active or enforced are updated
-        mock_checks_task.assert_not_called()
+        generate_agent_checks_from_policies_task.assert_not_called()
 
         data = {
             "name": "Test Policy Update",
@@ -118,7 +118,9 @@ class TestPolicyViews(TacticalTestCase):
 
         resp = self.client.put(url, data, format="json")
         self.assertEqual(resp.status_code, 200)
-        mock_checks_task.assert_called_with(policypk=policy.pk, create_tasks=True)
+        generate_agent_checks_from_policies_task.assert_called_with(
+            policypk=policy.pk, create_tasks=True
+        )
 
         self.check_not_authenticated("put", url)
 
@@ -478,8 +480,7 @@ class TestPolicyTasks(TacticalTestCase):
         # setup data
         policy = baker.make("automation.Policy", active=True)
         checks = self.create_checks(policy=policy)
-        site = baker.make("clients.Site")
-        agent = baker.make_recipe("agents.agent", site=site, policy=policy)
+        agent = baker.make_recipe("agents.agent", policy=policy)
 
         # test policy assigned to agent
         generate_agent_checks_from_policies_task(policy.id)
@@ -543,70 +544,241 @@ class TestPolicyTasks(TacticalTestCase):
             7,
         )
 
-    def test_generating_agent_policy_checks_by_location(self):
-        from .tasks import generate_agent_checks_by_location_task
+    @patch("automation.tasks.generate_agent_checks_by_location_task.delay")
+    def test_generating_agent_policy_checks_by_location(
+        self, generate_agent_checks_by_location_task
+    ):
+        from automation.tasks import (
+            generate_agent_checks_by_location_task as generate_agent_checks,
+        )
 
-        # this test is sort of obsolete since I added the policy logic to the client/site model save method
         # setup data
         policy = baker.make("automation.Policy", active=True)
         self.create_checks(policy=policy)
-        clients = baker.make(
-            "clients.Client",
-            _quantity=2,
-            server_policy=policy,
-            workstation_policy=policy,
-        )
-        sites = baker.make("clients.Site", client=cycle(clients), _quantity=4)
-        server_agent = baker.make_recipe("agents.server_agent", site=sites[0])
-        workstation_agent = baker.make_recipe("agents.workstation_agent", site=sites[2])
-        agent1 = baker.make_recipe("agents.server_agent", site=sites[1])
-        agent2 = baker.make_recipe("agents.workstation_agent", site=sites[3])
 
-        generate_agent_checks_by_location_task(
-            {"site_id": sites[0].id},
-            "server",
+        server_agent = baker.make_recipe("agents.server_agent")
+        workstation_agent = baker.make_recipe("agents.workstation_agent")
+
+        # no checks should be preset on agents
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 0)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
+
+        # set workstation policy on client and policy checks 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_checks(
+            location={"site__client_id": workstation_agent.client.pk},
+            mon_type="workstation",
             create_tasks=True,
         )
 
-        # server_agent should have policy checks and the other agents should not
-        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
+        # make sure the checks were added
         self.assertEqual(
             Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 7
         )
-        self.assertEqual(Agent.objects.get(pk=agent1.id).agentchecks.count(), 7)
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 0)
 
-        generate_agent_checks_by_location_task(
-            {"site__client_id": clients[0].id},
-            "workstation",
+        # 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,
         )
-        # workstation_agent should now have policy checks and the other agents should not
+        generate_agent_checks_by_location_task.reset_mock()
+
+        generate_agent_checks(
+            location={"site__client_id": workstation_agent.client.pk},
+            mon_type="workstation",
+            create_tasks=True,
+        )
+
+        # make sure the checks were removed
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 0)
+
+        # set server policy on client and policy checks 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_checks(
+            location={"site__client_id": server_agent.client.pk},
+            mon_type="server",
+            create_tasks=True,
+        )
+
+        # make sure checks were added
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
+
+        # 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_checks(
+            location={"site__client_id": server_agent.client.pk},
+            mon_type="server",
+            create_tasks=True,
+        )
+
+        # make sure checks were removed
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 0)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
+
+        # 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_checks(
+            location={"site_id": workstation_agent.site.pk},
+            mon_type="workstation",
+            create_tasks=True,
+        )
+
+        # make sure checks were added on workstation
         self.assertEqual(
             Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 7
         )
-        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
-        self.assertEqual(Agent.objects.get(pk=agent1.id).agentchecks.count(), 7)
-        self.assertEqual(Agent.objects.get(pk=agent2.id).agentchecks.count(), 7)
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 0)
 
-    def test_generating_policy_checks_for_all_agents(self):
-        from .tasks import generate_all_agent_checks_task
+        # 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_checks(
+            location={"site_id": workstation_agent.site.pk},
+            mon_type="workstation",
+            create_tasks=True,
+        )
+
+        # make sure checks were removed
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 0)
+
+        # 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_checks(
+            location={"site_id": server_agent.site.pk},
+            mon_type="server",
+            create_tasks=True,
+        )
+
+        # make sure checks were added
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 7)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
+
+        # 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_checks(
+            location={"site_id": server_agent.site.pk},
+            mon_type="server",
+            create_tasks=True,
+        )
+
+        # make sure checks were removed
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).agentchecks.count(), 0)
+        self.assertEqual(
+            Agent.objects.get(pk=workstation_agent.id).agentchecks.count(), 0
+        )
+
+    @patch("automation.tasks.generate_all_agent_checks_task.delay")
+    def test_generating_policy_checks_for_all_agents(
+        self, generate_all_agent_checks_task
+    ):
+        from .tasks import generate_all_agent_checks_task as generate_all_checks
         from core.models import CoreSettings
 
         # setup data
         policy = baker.make("automation.Policy", active=True)
         self.create_checks(policy=policy)
 
-        site = baker.make("clients.Site")
-        server_agents = baker.make_recipe("agents.server_agent", site=site, _quantity=3)
-        workstation_agents = baker.make_recipe(
-            "agents.workstation_agent", site=site, _quantity=4
-        )
+        server_agents = baker.make_recipe("agents.server_agent", _quantity=3)
+        workstation_agents = baker.make_recipe("agents.workstation_agent", _quantity=4)
         core = CoreSettings.objects.first()
         core.server_policy = policy
-        core.workstation_policy = policy
         core.save()
 
-        generate_all_agent_checks_task("server", create_tasks=True)
+        generate_all_agent_checks_task.assert_called_with(
+            mon_type="server", create_tasks=True
+        )
+        generate_all_agent_checks_task.reset_mock()
+        generate_all_checks(mon_type="server", create_tasks=True)
 
         # all servers should have 7 checks
         for agent in server_agents:
@@ -615,24 +787,50 @@ class TestPolicyTasks(TacticalTestCase):
         for agent in workstation_agents:
             self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 0)
 
-        generate_all_agent_checks_task("workstation", create_tasks=True)
+        core.server_policy = None
+        core.workstation_policy = policy
+        core.save()
 
-        # all agents should have 7 checks now
+        generate_all_agent_checks_task.assert_any_call(
+            mon_type="workstation", create_tasks=True
+        )
+        generate_all_agent_checks_task.assert_any_call(
+            mon_type="server", create_tasks=True
+        )
+        generate_all_agent_checks_task.reset_mock()
+        generate_all_checks(mon_type="server", create_tasks=True)
+        generate_all_checks(mon_type="workstation", create_tasks=True)
+
+        # all workstations should have 7 checks
         for agent in server_agents:
-            self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 7)
+            self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 0)
 
         for agent in workstation_agents:
             self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 7)
 
+        core.workstation_policy = None
+        core.save()
+
+        generate_all_agent_checks_task.assert_called_with(
+            mon_type="workstation", create_tasks=True
+        )
+        generate_all_agent_checks_task.reset_mock()
+        generate_all_checks(mon_type="workstation", create_tasks=True)
+
+        # nothing should have the checks
+        for agent in server_agents:
+            self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 0)
+
+        for agent in workstation_agents:
+            self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 0)
+
     def test_delete_policy_check(self):
         from .tasks import delete_policy_check_task
         from .models import Policy
 
         policy = baker.make("automation.Policy", active=True)
         self.create_checks(policy=policy)
-        site = baker.make("clients.Site")
-        agent = baker.make_recipe("agents.server_agent", site=site, policy=policy)
-        agent.generate_checks_from_policies()
+        agent = baker.make_recipe("agents.server_agent", policy=policy)
 
         # make sure agent has 7 checks
         self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 7)
@@ -657,7 +855,6 @@ class TestPolicyTasks(TacticalTestCase):
         policy = baker.make("automation.Policy", active=True)
         self.create_checks(policy=policy)
         agent = baker.make_recipe("agents.server_agent", policy=policy)
-        agent.generate_checks_from_policies()
 
         # make sure agent has 7 checks
         self.assertEqual(Agent.objects.get(pk=agent.id).agentchecks.count(), 7)
@@ -689,8 +886,7 @@ class TestPolicyTasks(TacticalTestCase):
         tasks = baker.make(
             "autotasks.AutomatedTask", policy=policy, name=seq("Task"), _quantity=3
         )
-        site = baker.make("clients.Site")
-        agent = baker.make_recipe("agents.server_agent", site=site, policy=policy)
+        agent = baker.make_recipe("agents.server_agent", policy=policy)
 
         generate_agent_tasks_from_policies_task(policy.id)
 
@@ -711,48 +907,205 @@ class TestPolicyTasks(TacticalTestCase):
                 self.assertEqual(task.parent_task, tasks[2].id)
                 self.assertEqual(task.name, tasks[2].name)
 
-    def test_generate_agent_tasks_by_location(self):
-        from .tasks import generate_agent_tasks_by_location_task
+    @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,
+        )
 
-        # test is sort of obsolete since I most the policy changes to the client/site save method
         # setup data
         policy = baker.make("automation.Policy", active=True)
         baker.make(
             "autotasks.AutomatedTask", policy=policy, name=seq("Task"), _quantity=3
         )
-        clients = baker.make(
-            "clients.Client",
-            _quantity=2,
-            server_policy=policy,
-            workstation_policy=policy,
+
+        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
         )
-        sites = baker.make("clients.Site", client=cycle(clients), _quantity=4)
-        server_agent = baker.make_recipe("agents.server_agent", site=sites[0])
-        workstation_agent = baker.make_recipe("agents.workstation_agent", site=sites[2])
-        agent1 = baker.make_recipe("agents.agent", site=sites[1])
-        agent2 = baker.make_recipe("agents.agent", site=sites[3])
+        self.assertEqual(Agent.objects.get(pk=server_agent.id).autotasks.count(), 0)
 
-        generate_agent_tasks_by_location_task({"site_id": sites[0].id}, "server")
+        # set workstation policy on client and policy tasks should be there
+        workstation_agent.client.workstation_policy = policy
+        workstation_agent.client.save()
 
-        # everything should have 3 tasks
+        # 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.site.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.site.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)
-        self.assertEqual(Agent.objects.get(pk=agent1.id).autotasks.count(), 3)
-        self.assertEqual(Agent.objects.get(pk=agent2.id).autotasks.count(), 3)
 
-        generate_agent_tasks_by_location_task(
-            {"site__client_id": clients[0].id}, "workstation"
+        # 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"
         )
 
-        # everything should have 3 tasks
+        # 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)
-        self.assertEqual(Agent.objects.get(pk=agent1.id).autotasks.count(), 3)
-        self.assertEqual(Agent.objects.get(pk=agent2.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):
@@ -760,13 +1113,13 @@ class TestPolicyTasks(TacticalTestCase):
 
         policy = baker.make("automation.Policy", active=True)
         tasks = baker.make("autotasks.AutomatedTask", policy=policy, _quantity=3)
-        site = baker.make("clients.Site")
-        agent = baker.make_recipe("agents.server_agent", site=site, policy=policy)
-        agent.generate_tasks_from_policies()
+        agent = baker.make_recipe("agents.server_agent", policy=policy)
 
         delete_policy_autotask_task(tasks[0].id)
 
-        delete_win_task_schedule.assert_called_with(agent.autotasks.first().id)
+        delete_win_task_schedule.assert_called_with(
+            agent.autotasks.get(parent_task=tasks[0].id).id
+        )
 
     @patch("autotasks.tasks.run_win_task.delay")
     def test_run_policy_task(self, run_win_task):
@@ -781,25 +1134,26 @@ class TestPolicyTasks(TacticalTestCase):
         for task in tasks:
             run_win_task.assert_any_call(task.id)
 
-    @patch("agents.models.Agent.nats_cmd")
-    def test_update_policy_tasks(self, nats_cmd):
+    @patch("autotasks.tasks.enable_or_disable_win_task.delay")
+    def test_update_policy_tasks(self, enable_or_disable_win_task):
         from .tasks import update_policy_task_fields_task
-        from autotasks.models import AutomatedTask
-
-        nats_cmd.return_value = "ok"
 
         # setup data
         policy = baker.make("automation.Policy", active=True)
         tasks = baker.make(
             "autotasks.AutomatedTask", enabled=True, policy=policy, _quantity=3
         )
-        site = baker.make("clients.Site")
-        agent = baker.make_recipe("agents.server_agent", site=site, policy=policy)
-        agent.generate_tasks_from_policies()
+        agent = baker.make_recipe("agents.server_agent", policy=policy)
 
         tasks[0].enabled = False
         tasks[0].save()
 
         update_policy_task_fields_task(tasks[0].id)
+        enable_or_disable_win_task.assert_not_called()
 
-        self.assertFalse(AutomatedTask.objects.get(parent_task=tasks[0].id).enabled)
+        self.assertFalse(agent.autotasks.get(parent_task=tasks[0].id).enabled)
+
+        update_policy_task_fields_task(tasks[0].id, update_agent=True)
+        enable_or_disable_win_task.assert_called_with(
+            agent.autotasks.get(parent_task=tasks[0].id).id, False
+        )

api/tacticalrmm/automation/views.py

@@ -161,7 +161,7 @@ class GetRelated(APIView):
         ).data
 
         response["agents"] = AgentHostnameSerializer(
-            policy.related_agents(),
+            policy.related_agents().only("pk", "hostname"),
             many=True,
         ).data
 

api/tacticalrmm/autotasks/tests.py

@@ -150,7 +150,9 @@ class TestAutotaskViews(TacticalTestCase):
 
         resp = self.client.patch(url, data, format="json")
         self.assertEqual(resp.status_code, 200)
-        update_policy_task_fields_task.assert_called_with(policy_task.id, True)
+        update_policy_task_fields_task.assert_called_with(
+            policy_task.id, update_agent=True
+        )
 
         self.check_not_authenticated("patch", url)
 

api/tacticalrmm/autotasks/views.py

@@ -107,7 +107,7 @@ class AutoTask(APIView):
                 enable_or_disable_win_task.delay(pk=task.pk, action=action)
 
             else:
-                update_policy_task_fields_task.delay(task.pk)
+                update_policy_task_fields_task.delay(task.pk, update_agent=True)
 
             task.enabled = action
             task.save(update_fields=["enabled"])

api/tacticalrmm/checks/baker_recipes.py

@@ -1,7 +1,6 @@
-from .models import Check
-from model_bakery.recipe import Recipe, seq
+from model_bakery.recipe import Recipe
 
-check = Recipe(Check)
+check = Recipe("checks.Check")
 
 diskspace_check = check.extend(
     check_type="diskspace", disk="C:", warning_threshold=30, error_threshold=75

api/tacticalrmm/core/models.py

@@ -93,7 +93,7 @@ class CoreSettings(BaseAuditModel):
                 pass
 
         old_settings = type(self).objects.get(pk=self.pk) if self.pk else None
-        super(CoreSettings, self).save(*args, **kwargs)
+        super(BaseAuditModel, self).save(*args, **kwargs)
 
         # check if server polcies have changed and initiate task to reapply policies if so
         if old_settings and old_settings.server_policy != self.server_policy:
@@ -101,7 +101,9 @@ class CoreSettings(BaseAuditModel):
 
         # check if workstation polcies have changed and initiate task to reapply policies if so
         if old_settings and old_settings.workstation_policy != self.workstation_policy:
-            generate_all_agent_checks_task.delay(mon_type="server", create_tasks=True)
+            generate_all_agent_checks_task.delay(
+                mon_type="workstation", create_tasks=True
+            )
 
     def __str__(self):
         return "Global Site Settings"

api/tacticalrmm/logs/baker_recipes.py

@@ -1,4 +1,3 @@
-from logs.models import AuditLog
 from model_bakery.recipe import Recipe
 from itertools import cycle
 
@@ -19,10 +18,10 @@ object_actions = ["add", "modify", "view", "delete"]
 agent_actions = ["remote_session", "execute_script", "execute_command"]
 login_actions = ["failed_login", "login"]
 
-agent_logs = Recipe(AuditLog, action=cycle(agent_actions), object_type="agent")
+agent_logs = Recipe("logs.AuditLog", action=cycle(agent_actions), object_type="agent")
 
 object_logs = Recipe(
-    AuditLog, action=cycle(object_actions), object_type=cycle(object_types)
+    "logs.AuditLog", action=cycle(object_actions), object_type=cycle(object_types)
 )
 
-login_logs = Recipe(AuditLog, action=cycle(login_actions), object_type="user")
+login_logs = Recipe("logs.AuditLog", action=cycle(login_actions), object_type="user")

api/tacticalrmm/logs/tests.py

@@ -147,8 +147,8 @@ class TestAuditViews(TacticalTestCase):
     def test_options_filter(self):
         url = "/logs/auditlogs/optionsfilter/"
 
-        baker.make("agents.Agent", hostname=seq("AgentHostname"), _quantity=5)
-        baker.make("agents.Agent", hostname=seq("Server"), _quantity=3)
+        baker.make_recipe("agents.agent", hostname=seq("AgentHostname"), _quantity=5)
+        baker.make_recipe("agents.agent", hostname=seq("Server"), _quantity=3)
         baker.make("accounts.User", username=seq("Username"), _quantity=7)
         baker.make("accounts.User", username=seq("soemthing"), _quantity=3)
 
@@ -194,7 +194,8 @@ class TestAuditViews(TacticalTestCase):
 
     def test_all_pending_actions(self):
         url = "/logs/allpendingactions/"
-        pending_actions = baker.make("logs.PendingAction", _quantity=6)
+        agent = baker.make_recipe("agents.agent")
+        pending_actions = baker.make("logs.PendingAction", agent=agent, _quantity=6)
 
         resp = self.client.get(url, format="json")
         serializer = PendingActionSerializer(pending_actions, many=True)

api/tacticalrmm/scripts/baker_recipes.py

@@ -1,8 +1,7 @@
-from .models import Script
 from model_bakery.recipe import Recipe
 
 script = Recipe(
-    Script,
+    "scripts.Script",
     name="Test Script",
     description="Test Desc",
     shell="cmd",

api/tacticalrmm/services/tests.py

@@ -8,6 +8,7 @@ from agents.models import Agent
 class TestServiceViews(TacticalTestCase):
     def setUp(self):
         self.authenticate()
+        self.setup_coresettings()
 
     def test_default_services(self):
         url = "/services/defaultservices/"

api/tacticalrmm/software/tests.py

@@ -8,6 +8,7 @@ from .models import ChocoLog
 class TestSoftwareViews(TacticalTestCase):
     def setUp(self):
         self.authenticate()
+        self.setup_coresettings()
 
     def test_chocos_get(self):
         url = "/software/chocos/"
@@ -63,6 +64,9 @@ class TestSoftwareViews(TacticalTestCase):
 
 
 class TestSoftwareTasks(TacticalTestCase):
+    def setUp(self):
+        self.setup_coresettings()
+
     @patch("agents.models.Agent.nats_cmd")
     def test_install_program(self, nats_cmd):
         from .tasks import install_program

api/tacticalrmm/winupdate/tests.py

@@ -88,7 +88,8 @@ class TestWinUpdateViews(TacticalTestCase):
 
     def test_edit_policy(self):
         url = "/winupdate/editpolicy/"
-        winupdate = baker.make("winupdate.WinUpdate")
+        agent = baker.make_recipe("agents.agent")
+        winupdate = baker.make("winupdate.WinUpdate", agent=agent)
 
         invalid_data = {"pk": 500, "policy": "inherit"}
         # test a call where winupdate doesn't exist

web/src/components/AlertsIcon.vue

@@ -69,7 +69,7 @@ export default {
         })
         .catch(e => {
           this.$q.loading.hide();
-          this.notifyError("Unable to get alerts");
+          //this.notifyError("Unable to get alerts");
         });
     },
     showOverview() {