From a65eb62a54c6fb2af8738bc8d2b24fa388a39634 Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Fri, 29 Jan 2021 00:34:18 +0000
Subject: [PATCH] checkrunner changes
 wh1te909/rmmagent@10a0935f1b75c0f43cbd1689f3f4f1f45a42beb6

---
 api/tacticalrmm/checks/tests.py   | 42 ++++++++++++++++++++++++++++++-
 api/tacticalrmm/checks/views.py   | 14 +++++++++--
 web/src/components/AgentTable.vue | 13 +++++++---
 3 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/api/tacticalrmm/checks/tests.py b/api/tacticalrmm/checks/tests.py
index 5b43d06f..64fbc164 100644
--- a/api/tacticalrmm/checks/tests.py
+++ b/api/tacticalrmm/checks/tests.py
@@ -2,9 +2,9 @@ from checks.models import CheckHistory
 from tacticalrmm.test import TacticalTestCase
 from .serializers import CheckSerializer
 from django.utils import timezone as djangotime
+from unittest.mock import patch
 
 from model_bakery import baker
-from itertools import cycle
 
 
 class TestCheckViews(TacticalTestCase):
@@ -184,6 +184,46 @@ class TestCheckViews(TacticalTestCase):
 
         self.check_not_authenticated("patch", url_a)
 
+    @patch("agents.models.Agent.nats_cmd")
+    def test_run_checks(self, nats_cmd):
+        agent = baker.make_recipe("agents.agent", version="1.4.1")
+        agent_old = baker.make_recipe("agents.agent", version="1.0.2")
+        agent_b4_141 = baker.make_recipe("agents.agent", version="1.4.0")
+
+        url = f"/checks/runchecks/{agent_old.pk}/"
+        r = self.client.get(url)
+        self.assertEqual(r.status_code, 400)
+        self.assertEqual(r.json(), "Requires agent version 1.1.0 or greater")
+
+        url = f"/checks/runchecks/{agent_b4_141.pk}/"
+        r = self.client.get(url)
+        self.assertEqual(r.status_code, 200)
+        nats_cmd.assert_called_with({"func": "runchecks"}, wait=False)
+
+        nats_cmd.reset_mock()
+        nats_cmd.return_value = "busy"
+        url = f"/checks/runchecks/{agent.pk}/"
+        r = self.client.get(url)
+        self.assertEqual(r.status_code, 400)
+        nats_cmd.assert_called_with({"func": "runchecks"}, timeout=15)
+        self.assertEqual(r.json(), f"Checks are already running on {agent.hostname}")
+
+        nats_cmd.reset_mock()
+        nats_cmd.return_value = "ok"
+        url = f"/checks/runchecks/{agent.pk}/"
+        r = self.client.get(url)
+        self.assertEqual(r.status_code, 200)
+        nats_cmd.assert_called_with({"func": "runchecks"}, timeout=15)
+        self.assertEqual(r.json(), f"Checks will now be re-run on {agent.hostname}")
+
+        nats_cmd.reset_mock()
+        nats_cmd.return_value = "timeout"
+        url = f"/checks/runchecks/{agent.pk}/"
+        r = self.client.get(url)
+        self.assertEqual(r.status_code, 400)
+        nats_cmd.assert_called_with({"func": "runchecks"}, timeout=15)
+        self.assertEqual(r.json(), "Unable to contact the agent")
+
     def test_get_check_history(self):
         # setup data
         agent = baker.make_recipe("agents.agent")
diff --git a/api/tacticalrmm/checks/views.py b/api/tacticalrmm/checks/views.py
index 2f601a8c..3900d9bc 100644
--- a/api/tacticalrmm/checks/views.py
+++ b/api/tacticalrmm/checks/views.py
@@ -1,4 +1,5 @@
 import asyncio
+from packaging import version as pyver
 
 from django.shortcuts import get_object_or_404
 from django.db.models import Q
@@ -168,8 +169,17 @@ def run_checks(request, pk):
     if not agent.has_nats:
         return notify_error("Requires agent version 1.1.0 or greater")
 
-    asyncio.run(agent.nats_cmd({"func": "runchecks"}, wait=False))
-    return Response(agent.hostname)
+    if pyver.parse(agent.version) >= pyver.parse("1.4.1"):
+        r = asyncio.run(agent.nats_cmd({"func": "runchecks"}, timeout=15))
+        if r == "busy":
+            return notify_error(f"Checks are already running on {agent.hostname}")
+        elif r == "ok":
+            return Response(f"Checks will now be re-run on {agent.hostname}")
+        else:
+            return notify_error("Unable to contact the agent")
+    else:
+        asyncio.run(agent.nats_cmd({"func": "runchecks"}, wait=False))
+        return Response(f"Checks will now be re-run on {agent.hostname}")
 
 
 @api_view()
diff --git a/web/src/components/AgentTable.vue b/web/src/components/AgentTable.vue
index 0b83d7f9..cbcc2e2f 100644
--- a/web/src/components/AgentTable.vue
+++ b/web/src/components/AgentTable.vue
@@ -541,10 +541,17 @@ export default {
       window.open(url, "", "scrollbars=no,location=no,status=no,toolbar=no,menubar=no,width=1280,height=826");
     },
     runChecks(pk) {
-      axios
+      this.$q.loading.show();
+      this.$axios
         .get(`/checks/runchecks/${pk}/`)
-        .then(r => this.notifySuccess(`Checks will now be re-run on ${r.data}`))
-        .catch(e => this.notifyError(e.response.data));
+        .then(r => {
+          this.$q.loading.hide();
+          this.notifySuccess(r.data);
+        })
+        .catch(e => {
+          this.$q.loading.hide();
+          this.notifyError(e.response.data);
+        });
     },
     removeAgent(pk, name) {
       this.$q