diff --git a/api/tacticalrmm/agents/tests.py b/api/tacticalrmm/agents/tests.py index 1a8b5786..9213df26 100644 --- a/api/tacticalrmm/agents/tests.py +++ b/api/tacticalrmm/agents/tests.py @@ -187,20 +187,17 @@ class TestAgentViews(TacticalTestCase): self.check_not_authenticated("get", url) - @patch("agents.models.Agent.salt_api_cmd") - def test_power_action(self, mock_ret): + @patch("agents.models.Agent.nats_cmd") + def test_power_action(self, nats_cmd): url = f"/agents/poweraction/" data = {"pk": self.agent.pk, "action": "rebootnow"} - mock_ret.return_value = True + nats_cmd.return_value = "ok" r = self.client.post(url, data, format="json") self.assertEqual(r.status_code, 200) + nats_cmd.assert_called_with({"func": "rebootnow"}, timeout=10) - mock_ret.return_value = "error" - r = self.client.post(url, data, format="json") - self.assertEqual(r.status_code, 400) - - mock_ret.return_value = False + nats_cmd.return_value = "timeout" r = self.client.post(url, data, format="json") self.assertEqual(r.status_code, 400) diff --git a/api/tacticalrmm/agents/views.py b/api/tacticalrmm/agents/views.py index df6bc472..9c8bd3ac 100644 --- a/api/tacticalrmm/agents/views.py +++ b/api/tacticalrmm/agents/views.py @@ -203,19 +203,13 @@ def get_event_log(request, pk, logtype, days): @api_view(["POST"]) def power_action(request): - pk = request.data["pk"] - action = request.data["action"] - agent = get_object_or_404(Agent, pk=pk) - if action == "rebootnow": - logger.info(f"{agent.hostname} was scheduled for immediate reboot") - r = agent.salt_api_cmd( - timeout=30, - func="system.reboot", - arg=3, - kwargs={"in_seconds": True}, - ) - if r == "timeout" or r == "error" or (isinstance(r, bool) and not r): - return notify_error("Unable to contact the agent") + agent = get_object_or_404(Agent, pk=request.data["pk"]) + if not agent.has_nats: + return notify_error("Requires agent version 1.1.0 or greater") + if request.data["action"] == "rebootnow": + r = asyncio.run(agent.nats_cmd({"func": "rebootnow"}, timeout=10)) + if r != "ok": + return notify_error("Unable to contact the agent") return Response("ok")